Repository: puzpuzpuz/xsync
Pull request #68 (2022-10-29 10:30:03)

Load and compute
Decided to give it a try.  This is a PR to go into your PR here: https://github.com/puzpuzpuz/xsync/pull/65

puzpuzpuz: I'll handle the pending comment in a separate PR. Thanks for the contribution!

Patch:
diff --git a/mapof.go b/mapof.go
index 32d1cc9..025df82 100644
--- a/mapof.go
+++ b/mapof.go
@@ -112,10 +112,11 @@ func (m *MapOf[K, V]) Load(key K) (value V, ok bool) {
 func (m *MapOf[K, V]) Store(key K, value V) {
 	m.doStore(
 		key,
-		func() V {
-			return value
+		func(V, bool) (V, bool) {
+			return value, false
 		},
 		false,
+		false,
 	)
 }
 
@@ -125,10 +126,11 @@ func (m *MapOf[K, V]) Store(key K, value V) {
 func (m *MapOf[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 	return m.doStore(
 		key,
-		func() V {
-			return value
+		func(V, bool) (V, bool) {
+			return value, false
 		},
 		true,
+		false,
 	)
 }
 
@@ -140,10 +142,11 @@ func (m *MapOf[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 func (m *MapOf[K, V]) LoadAndStore(key K, value V) (actual V, loaded bool) {
 	return m.doStore(
 		key,
-		func() V {
-			return value
+		func(V, bool) (V, bool) {
+			return value, false
 		},
 		false,
+		false,
 	)
 }
 
@@ -152,14 +155,21 @@ func (m *MapOf[K, V]) LoadAndStore(key K, value V) (actual V, loaded bool) {
 // returns the computed value. The loaded result is true if the value
 // was loaded, false if stored.
 func (m *MapOf[K, V]) LoadOrCompute(key K, valueFn func() V) (actual V, loaded bool) {
-	return m.doStore(key, valueFn, true)
+	return m.doStore(
+		key,
+		func(V, bool) (V, bool) {
+			return valueFn(), false
+		},
+		true,
+		false,
+	)
 }
 
-func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bool) {
+func (m *MapOf[K, V]) doStore(key K, valueFn func(oldValue V, loaded bool) (V, bool), loadIfExists, computedOk bool) (V, bool) {
 	// Read-only path.
 	if loadIfExists {
 		if v, ok := m.Load(key); ok {
-			return v, true
+			return v, !computedOk
 		}
 	}
 	// Write path.
@@ -169,6 +179,7 @@ func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bo
 			emptyb   *bucketPadded
 			emptyidx int
 		)
+		hintNonEmpty := 0
 		table := (*mapTable)(atomic.LoadPointer(&m.table))
 		tableLen := len(table.buckets)
 		hash := m.hasher(table.seed, key)
@@ -204,13 +215,34 @@ func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bo
 					vp := b.values[i]
 					if loadIfExists {
 						unlockBucket(&rootb.topHashMutex)
-						return derefTypedValue[V](vp), true
+						return derefTypedValue[V](vp), !computedOk
 					}
 					// In-place update case. We get a copy of the value via an
 					// interface{} on each call, thus the live value pointers are
 					// unique. Otherwise atomic snapshot won't be correct in case
 					// of multiple Store calls using the same value.
-					value := valueFn()
+					value, del := valueFn(derefTypedValue[V](vp), true)
+					if del {
+						// Deletion case. First we update the value, then the key.
+						// This is important for atomic snapshot states.
+						atomic.StoreUint64(&b.topHashMutex, eraseTopHash(topHashes, i))
+						atomic.StorePointer(&b.values[i], nil)
+						atomic.StorePointer(&b.keys[i], nil)
+						leftEmpty := false
+						if hintNonEmpty == 0 {
+							leftEmpty = isEmpty(b)
+						}
+						unlockBucket(&rootb.topHashMutex)
+						table.addSize(bidx, -1)
+						// Might need to shrink the table.
+						if leftEmpty {
+							m.resize(table, mapShrinkHint)
+						}
+						if computedOk {
+							return value, false
+						}
+						return derefTypedValue[V](vp), true
+					}
 					var wv interface{} = value
 					nvp := unsafe.Pointer(&wv)
 					if assertionsEnabled && vp == nvp {
@@ -219,22 +251,31 @@ func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bo
 					atomic.StorePointer(&b.values[i], nvp)
 					unlockBucket(&rootb.topHashMutex)
 					// LoadAndStore expects the old value to be returned.
+					if computedOk {
+						return value, true
+					}
 					return derefTypedValue[V](vp), true
 				}
+				hintNonEmpty++
 			}
 			if b.next == nil {
+				var zeroedV V
+				value, del := valueFn(zeroedV, false)
+				if del {
+					unlockBucket(&rootb.topHashMutex)
+					return value, false
+				}
+				var wv interface{} = value
 				if emptyb != nil {
 					// Insertion case. First we update the value, then the key.
 					// This is important for atomic snapshot states.
 					topHashes = atomic.LoadUint64(&emptyb.topHashMutex)
 					atomic.StoreUint64(&emptyb.topHashMutex, storeTopHash(hash, topHashes, emptyidx))
-					value := valueFn()
-					var wv interface{} = value
 					atomic.StorePointer(&emptyb.values[emptyidx], unsafe.Pointer(&wv))
 					atomic.StorePointer(&emptyb.keys[emptyidx], unsafe.Pointer(&key))
 					unlockBucket(&rootb.topHashMutex)
 					table.addSize(bidx, 1)
-					return value, false
+					return value, computedOk
 				}
 				growThreshold := float64(tableLen) * entriesPerMapBucket * mapLoadFactor
 				if table.sumSize() > int64(growThreshold) {
@@ -246,20 +287,29 @@ func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bo
 				// Create and append a new bucket.
 				newb := new(bucketPadded)
 				newb.keys[0] = unsafe.Pointer(&key)
-				value := valueFn()
-				var wv interface{} = value
 				newb.values[0] = unsafe.Pointer(&wv)
 				newb.topHashMutex = storeTopHash(hash, newb.topHashMutex, 0)
 				atomic.StorePointer(&b.next, unsafe.Pointer(newb))
 				unlockBucket(&rootb.topHashMutex)
 				table.addSize(bidx, 1)
-				return value, false
+				return value, computedOk
 			}
 			b = (*bucketPadded)(b.next)
 		}
 	}
 }
 
+// Compute either sets the computed new value for the key or deletes
+// the value for the key. When the delete result of the valueFn function
+// is set to true, the value will be deleted, if it exists. When delete
+// is set to false, the value is updated to the newValue.
+// The ok result indicates whether value was computed and stored, thus, is
+// present in the map. The actual result contains the new value in cases where
+// the value was computed and stored. See the example for a few use cases.
+func (m *MapOf[K, V]) Compute(key K, valueFn func(oldValue V, loaded bool) (newValue V, delete bool)) (actual V, ok bool) {
+	return m.doStore(key, valueFn, false, true)
+}
+
 func (m *MapOf[K, V]) newerTableExists(table *mapTable) bool {
 	curTablePtr := atomic.LoadPointer(&m.table)
 	return uintptr(curTablePtr) != uintptr(unsafe.Pointer(table))


Test patch:
diff --git a/example_test.go b/example_test.go
index 68cbe35..59a21d8 100644
--- a/example_test.go
+++ b/example_test.go
@@ -5,6 +5,7 @@ package xsync_test
 
 import (
 	"encoding/binary"
+	"fmt"
 	"hash/maphash"
 	"time"
 
@@ -33,3 +34,53 @@ func ExampleNewTypedMapOf() {
 	age.Store(Person{"Ada", "Lovelace", 1815}, Y-1815)
 	age.Store(Person{"Charles", "Babbage", 1791}, Y-1791)
 }
+
+func ExampleMapOf_Compute() {
+	counts := xsync.NewIntegerMapOf[int, int]()
+
+	// Store a new value.
+	v, ok := counts.Compute(42, func(oldValue int, loaded bool) (newValue int, delete bool) {
+		// loaded is false here.
+		newValue = 42
+		delete = false
+		return
+	})
+	// v: 42, ok: true
+	fmt.Printf("v: %v, ok: %v\n", v, ok)
+
+	// Update an existing value.
+	v, ok = counts.Compute(42, func(oldValue int, loaded bool) (newValue int, delete bool) {
+		// loaded is true here.
+		newValue = oldValue + 42
+		delete = false
+		return
+	})
+	// v: 84, ok: true
+	fmt.Printf("v: %v, ok: %v\n", v, ok)
+
+	// Compare existing value, and set new value or keep old value, and store the oldVal
+	var oldVal int
+	minVal := 63
+	v, ok = counts.Compute(42, func(oldValue int, loaded bool) (newValue int, delete bool) {
+		oldVal = oldValue
+		if !loaded || oldValue < minVal {
+			newValue = minVal
+			delete = false
+			return
+		}
+		newValue = oldValue
+		delete = false
+		return
+	})
+	// v: 84, ok: true, oldVal: 84
+	fmt.Printf("v: %v, ok: %v, oldVal: %v\n", v, ok, oldVal)
+
+	// Delete an existing value.
+	v, ok = counts.Compute(42, func(oldValue int, loaded bool) (newValue int, delete bool) {
+		// loaded is true here.
+		delete = true
+		return
+	})
+	// v: 84, ok: false
+	fmt.Printf("v: %v, ok: %v\n", v, ok)
+}
diff --git a/map_test.go b/map_test.go
index 3cddbc2..e586096 100644
--- a/map_test.go
+++ b/map_test.go
@@ -337,6 +337,23 @@ func TestMapSerialLoadOrCompute(t *testing.T) {
 	}
 }
 
+func TestMapSerialLoadOrCompute_FunctionCalledOnce(t *testing.T) {
+	m := NewMap()
+	for i := 0; i < 100; {
+		m.LoadOrCompute(strconv.Itoa(i), func() (v interface{}) {
+			v, i = i, i+1
+			return v
+		})
+	}
+
+	m.Range(func(k string, v interface{}) bool {
+		if vi, ok := v.(int); !ok || strconv.Itoa(vi) != k {
+			t.Errorf("%sth key is not equal to value %d", k, v)
+		}
+		return true
+	})
+}
+
 func TestMapSerialStoreThenDelete(t *testing.T) {
 	const numEntries = 1000
 	m := NewMap()
diff --git a/mapof_test.go b/mapof_test.go
index 36b3b65..05b4804 100644
--- a/mapof_test.go
+++ b/mapof_test.go
@@ -395,19 +395,22 @@ func TestMapOfSerialLoadOrCompute(t *testing.T) {
 			t.Errorf("values do not match for %d: %v", i, v)
 		}
 	}
+}
 
-	t.Run("valueFn func should be called only once", func(t *testing.T) {
-		m := NewIntegerMapOf[int, int]()
-		for i := 0; i < 100; {
-			m.LoadOrCompute(i, func() (v int) { v, i = i, i+1; return v })
-		}
-
-		m.Range(func(k, v int) bool {
-			if k != v {
-				t.Fatalf("%dth key is not equal to value %d", k, v)
-			}
-			return true
+func TestMapOfSerialLoadOrCompute_FunctionCalledOnce(t *testing.T) {
+	m := NewIntegerMapOf[int, int]()
+	for i := 0; i < 100; {
+		m.LoadOrCompute(i, func() (v int) {
+			v, i = i, i+1
+			return v
 		})
+	}
+
+	m.Range(func(k, v int) bool {
+		if k != v {
+			t.Errorf("%dth key is not equal to value %d", k, v)
+		}
+		return true
 	})
 }