Repository: puzpuzpuz/xsync
Pull request #50 (2022-10-22 10:14:49)

Add LoadOrCompute method to Map and MapOf
This method is similar to `Map#computeIfAbsent()` in Java.

codecov-commenter: # [Codecov](https://codecov.io/gh/puzpuzpuz/xsync/pull/50?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov) Report
Base: **95.26**% // Head: **95.01**% // Decreases project coverage by **`-0.25%`** :warning:
> Coverage data is based on head [(`7ced804`)](https://codecov.io/gh/puzpuzpuz/xsync/pull/50?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov) compared to base [(`e5e826b`)](https://codecov.io/gh/puzpuzpuz/xsync/commit/e5e826b606afcc6a69a0e1b33211c34cb7b98128?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov).
> Patch coverage: 100.00% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main      #50      +/-   ##
==========================================
- Coverage   95.26%   95.01%   -0.26%     
==========================================
  Files           6        6              
  Lines         866      902      +36     
==========================================
+ Hits          825      857      +32     
- Misses         33       36       +3     
- Partials        8        9       +1     
```


| [Impacted Files](https://codecov.io/gh/puzpuzpuz/xsync/pull/50?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov) | Coverage Δ | |
|---|---|---|
| [map.go](https://codecov.io/gh/puzpuzpuz/xsync/pull/50/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov#diff-bWFwLmdv) | `94.85% <100.00%> (+0.22%)` | :arrow_up: |
| [mapof.go](https://codecov.io/gh/puzpuzpuz/xsync/pull/50/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov#diff-bWFwb2YuZ28=) | `93.49% <100.00%> (-0.93%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/puzpuzpuz/xsync/pull/50?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Andrey+Pechkurov).


Patch:
diff --git a/map.go b/map.go
index 7ef9993..4ff2d6e 100644
--- a/map.go
+++ b/map.go
@@ -167,25 +167,52 @@ func (m *Map) Load(key string) (value interface{}, ok bool) {
 
 // Store sets the value for a key.
 func (m *Map) Store(key string, value interface{}) {
-	m.doStore(key, value, false)
+	m.doStore(
+		key,
+		func() interface{} {
+			return value
+		},
+		false,
+	)
 }
 
 // LoadOrStore returns the existing value for the key if present.
 // Otherwise, it stores and returns the given value.
 // The loaded result is true if the value was loaded, false if stored.
 func (m *Map) LoadOrStore(key string, value interface{}) (actual interface{}, loaded bool) {
-	return m.doStore(key, value, true)
+	return m.doStore(
+		key,
+		func() interface{} {
+			return value
+		},
+		true,
+	)
 }
 
 // LoadAndStore returns the existing value for the key if present,
 // while setting the new value for the key.
-// Otherwise, it stores and returns the given value.
-// The loaded result is true if the value was loaded, false otherwise.
+// It stores the new value and returns the existing one, if present.
+// The loaded result is true if the existing value was loaded,
+// false otherwise.
 func (m *Map) LoadAndStore(key string, value interface{}) (actual interface{}, loaded bool) {
-	return m.doStore(key, value, false)
+	return m.doStore(
+		key,
+		func() interface{} {
+			return value
+		},
+		false,
+	)
+}
+
+// LoadOrCompute returns the existing value for the key if present.
+// Otherwise, it computes the value using the provided function and
+// returns the computed value. The loaded result is true if the value
+// was loaded, false if stored.
+func (m *Map) LoadOrCompute(key string, valueFn func() interface{}) (actual interface{}, loaded bool) {
+	return m.doStore(key, valueFn, true)
 }
 
-func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interface{}, bool) {
+func (m *Map) doStore(key string, valueFn func() interface{}, loadIfExists bool) (interface{}, bool) {
 	// Read-only path.
 	if loadIfExists {
 		if v, ok := m.Load(key); ok {
@@ -241,12 +268,14 @@ func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interfa
 					// 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()
 					nvp := unsafe.Pointer(&value)
 					if assertionsEnabled && vp == nvp {
 						panic("non-unique value pointer")
 					}
 					atomic.StorePointer(&b.values[i], nvp)
 					unlockBucket(&rootb.topHashMutex)
+					// LoadAndStore expects the old value to be returned.
 					return derefValue(vp), true
 				}
 			}
@@ -255,6 +284,7 @@ func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interfa
 					// Insertion case. First we update the value, then the key.
 					// This is important for atomic snapshot states.
 					atomic.StoreUint64(&b.topHashMutex, storeTopHash(hash, topHashes, emptyidx))
+					value := valueFn()
 					atomic.StorePointer(emptyvp, unsafe.Pointer(&value))
 					atomic.StorePointer(emptykp, unsafe.Pointer(&key))
 					unlockBucket(&rootb.topHashMutex)
@@ -271,6 +301,7 @@ func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interfa
 				// Create and append a new bucket.
 				newb := new(bucketPadded)
 				newb.keys[0] = unsafe.Pointer(&key)
+				value := valueFn()
 				newb.values[0] = unsafe.Pointer(&value)
 				newb.topHashMutex = storeTopHash(hash, topHashes, emptyidx)
 				atomic.StorePointer(&b.next, unsafe.Pointer(newb))
diff --git a/mapof.go b/mapof.go
index d073dd9..9c6d3cc 100644
--- a/mapof.go
+++ b/mapof.go
@@ -114,25 +114,52 @@ func (m *MapOf[K, V]) Load(key K) (value V, ok bool) {
 
 // Store sets the value for a key.
 func (m *MapOf[K, V]) Store(key K, value V) {
-	m.doStore(key, value, false)
+	m.doStore(
+		key,
+		func() V {
+			return value
+		},
+		false,
+	)
 }
 
 // LoadOrStore returns the existing value for the key if present.
 // Otherwise, it stores and returns the given value.
 // The loaded result is true if the value was loaded, false if stored.
 func (m *MapOf[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
-	return m.doStore(key, value, true)
+	return m.doStore(
+		key,
+		func() V {
+			return value
+		},
+		true,
+	)
 }
 
 // LoadAndStore returns the existing value for the key if present,
 // while setting the new value for the key.
-// Otherwise, it stores and returns the given value.
-// The loaded result is true if the value was loaded, false otherwise.
+// It stores the new value and returns the existing one, if present.
+// The loaded result is true if the existing value was loaded,
+// false otherwise.
 func (m *MapOf[K, V]) LoadAndStore(key K, value V) (actual V, loaded bool) {
-	return m.doStore(key, value, false)
+	return m.doStore(
+		key,
+		func() V {
+			return value
+		},
+		false,
+	)
+}
+
+// LoadOrCompute returns the existing value for the key if present.
+// Otherwise, it computes the value using the provided function and
+// 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)
 }
 
-func (m *MapOf[K, V]) doStore(key K, value V, loadIfExists bool) (V, bool) {
+func (m *MapOf[K, V]) doStore(key K, valueFn func() V, loadIfExists bool) (V, bool) {
 	// Read-only path.
 	if loadIfExists {
 		if v, ok := m.Load(key); ok {
@@ -188,6 +215,7 @@ func (m *MapOf[K, V]) doStore(key K, value V, loadIfExists bool) (V, bool) {
 					// 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()
 					var wv interface{} = value
 					nvp := unsafe.Pointer(&wv)
 					if assertionsEnabled && vp == nvp {
@@ -195,6 +223,7 @@ func (m *MapOf[K, V]) doStore(key K, value V, loadIfExists bool) (V, bool) {
 					}
 					atomic.StorePointer(&b.values[i], nvp)
 					unlockBucket(&rootb.topHashMutex)
+					// LoadAndStore expects the old value to be returned.
 					return derefTypedValue[V](vp), true
 				}
 			}
@@ -203,7 +232,8 @@ func (m *MapOf[K, V]) doStore(key K, value V, loadIfExists bool) (V, bool) {
 					// Insertion case. First we update the value, then the key.
 					// This is important for atomic snapshot states.
 					atomic.StoreUint64(&b.topHashMutex, storeTopHash(hash, topHashes, emptyidx))
-					var wv interface{} = value
+					value := valueFn()
+					var wv interface{} = valueFn()
 					atomic.StorePointer(emptyvp, unsafe.Pointer(&wv))
 					atomic.StorePointer(emptykp, unsafe.Pointer(&key))
 					unlockBucket(&rootb.topHashMutex)
@@ -220,6 +250,7 @@ func (m *MapOf[K, V]) doStore(key K, value V, loadIfExists bool) (V, bool) {
 				// 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, topHashes, emptyidx)


Test patch:
diff --git a/map_test.go b/map_test.go
index 3175faa..0312f95 100644
--- a/map_test.go
+++ b/map_test.go
@@ -169,6 +169,14 @@ func TestMapLoadOrStore_NonNilValue(t *testing.T) {
 	if v != newv {
 		t.Errorf("value does not match: %v", v)
 	}
+	newv2 := &foo{}
+	v, loaded = m.LoadOrStore("foo", newv2)
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != newv {
+		t.Errorf("value does not match: %v", v)
+	}
 }
 
 func TestMapLoadAndStore_NilValue(t *testing.T) {
@@ -307,6 +315,33 @@ func TestMapSerialLoadOrStore(t *testing.T) {
 	}
 }
 
+func TestMapSerialLoadOrCompute(t *testing.T) {
+	const numEntries = 1000
+	m := NewMap()
+	for i := 0; i < numEntries; i++ {
+		v, loaded := m.LoadOrCompute(strconv.Itoa(i), func() interface{} {
+			return i
+		})
+		if loaded {
+			t.Errorf("value not computed for %d", i)
+		}
+		if vi, ok := v.(int); ok && vi != i {
+			t.Errorf("values do not match for %d: %v", i, v)
+		}
+	}
+	for i := 0; i < numEntries; i++ {
+		v, loaded := m.LoadOrCompute(strconv.Itoa(i), func() interface{} {
+			return i
+		})
+		if !loaded {
+			t.Errorf("value not loaded for %d", i)
+		}
+		if vi, ok := v.(int); ok && vi != i {
+			t.Errorf("values do not match for %d: %v", i, v)
+		}
+	}
+}
+
 func TestMapSerialStoreThenDelete(t *testing.T) {
 	const numEntries = 1000
 	m := NewMap()
diff --git a/mapof_test.go b/mapof_test.go
index e965099..23f35ab 100644
--- a/mapof_test.go
+++ b/mapof_test.go
@@ -130,6 +130,14 @@ func TestMapOfLoadOrStore_NonNilValue(t *testing.T) {
 	if v != newv {
 		t.Errorf("value does not match: %v", v)
 	}
+	newv2 := &foo{}
+	v, loaded = m.LoadOrStore("foo", newv2)
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != newv {
+		t.Errorf("value does not match: %v", v)
+	}
 }
 
 func TestMapOfLoadAndStore_NilValue(t *testing.T) {
@@ -354,6 +362,33 @@ func TestMapOfSerialLoadOrStore(t *testing.T) {
 	}
 }
 
+func TestMapOfSerialLoadOrCompute(t *testing.T) {
+	const numEntries = 1000
+	m := NewMapOf[int]()
+	for i := 0; i < numEntries; i++ {
+		v, loaded := m.LoadOrCompute(strconv.Itoa(i), func() int {
+			return i
+		})
+		if loaded {
+			t.Errorf("value not computed for %d", i)
+		}
+		if v != i {
+			t.Errorf("values do not match for %d: %v", i, v)
+		}
+	}
+	for i := 0; i < numEntries; i++ {
+		v, loaded := m.LoadOrCompute(strconv.Itoa(i), func() int {
+			return i
+		})
+		if !loaded {
+			t.Errorf("value not loaded for %d", i)
+		}
+		if v != i {
+			t.Errorf("values do not match for %d: %v", i, v)
+		}
+	}
+}
+
 func TestMapOfSerialStoreThenDelete(t *testing.T) {
 	const numEntries = 1000
 	m := NewMapOf[int]()