Repository: puzpuzpuz/xsync
Pull request #41 (2022-07-27 15:07:25)

Add Map.LoadAndStore
Thanks for such a useful package

In practice, it is common to store new values while getting old values, similar to `atomic.Value.Swap()`, so this feature is added.

**Requirements:**

```go
var (
	a      atomic.Value
	i      int
	loaded bool
)

i = 1
v := a.Swap(i)
loaded = v != nil
if !loaded {
	v = 1
}
fmt.Println(v, loaded) // 1, false

i = 2
v = a.Swap(i)
loaded = v != nil
fmt.Println(v, loaded) // 1, true

i = 3
v = a.Swap(i)
loaded = v != nil
fmt.Println(v, loaded) // 2, true
```

**Implementation:**

```go
m := NewMapOf[int]()
i := 1
v, loaded := m.LoadAndStore("foo", i)
fmt.Println(v, loaded) // 1, false

i = 2
v, loaded = m.LoadAndStore("foo", i)
fmt.Println(v, loaded) // 1, true

i = 3
v, loaded = m.LoadAndStore("foo", i)
fmt.Println(v, loaded) // 2, true
```

If you find it useful, please merge, otherwise please close.


Patch:
diff --git a/map.go b/map.go
index f57930c..c1362b4 100644
--- a/map.go
+++ b/map.go
@@ -160,6 +160,14 @@ func (m *Map) LoadOrStore(key string, value interface{}) (actual interface{}, lo
 	return m.doStore(key, 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.
+func (m *Map) LoadAndStore(key string, value interface{}) (actual interface{}, loaded bool) {
+	return m.doStore(key, value, false)
+}
+
 func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interface{}, bool) {
 	if value == nil {
 		value = nilVal
@@ -205,8 +213,8 @@ func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interfa
 				continue
 			}
 			if key == derefKey(b.keys[i]) {
+				vp := b.values[i]
 				if loadIfExists {
-					vp := b.values[i]
 					b.mu.Unlock()
 					return derefValue(vp), true
 				}
@@ -216,7 +224,7 @@ func (m *Map) doStore(key string, value interface{}, loadIfExists bool) (interfa
 				// of multiple Store calls using the same value.
 				atomic.StorePointer(&b.values[i], unsafe.Pointer(&value))
 				b.mu.Unlock()
-				return value, false
+				return derefValue(vp), true
 			}
 		}
 		if emptykp != nil {
diff --git a/mapof.go b/mapof.go
index 6d3e83c..096c7e3 100644
--- a/mapof.go
+++ b/mapof.go
@@ -91,6 +91,14 @@ func (m *MapOf[V]) LoadOrStore(key string, value V) (actual V, loaded bool) {
 	return m.doStore(key, 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.
+func (m *MapOf[V]) LoadAndStore(key string, value V) (actual V, loaded bool) {
+	return m.doStore(key, value, false)
+}
+
 func (m *MapOf[V]) doStore(key string, value V, loadIfExists bool) (V, bool) {
 	// Read-only path.
 	if loadIfExists {
@@ -133,8 +141,8 @@ func (m *MapOf[V]) doStore(key string, value V, loadIfExists bool) (V, bool) {
 				continue
 			}
 			if key == derefKey(b.keys[i]) {
+				vp := b.values[i]
 				if loadIfExists {
-					vp := b.values[i]
 					b.mu.Unlock()
 					return derefTypedValue[V](vp), true
 				}
@@ -144,7 +152,7 @@ func (m *MapOf[V]) doStore(key string, value V, loadIfExists bool) (V, bool) {
 				// of multiple Store calls using the same value.
 				atomic.StorePointer(&b.values[i], unsafe.Pointer(&value))
 				b.mu.Unlock()
-				return value, false
+				return derefTypedValue[V](vp), true
 			}
 		}
 		if emptykp != nil {


Test patch:
diff --git a/map_test.go b/map_test.go
index a19c475..30e7ca8 100644
--- a/map_test.go
+++ b/map_test.go
@@ -113,6 +113,53 @@ func TestMapLoadOrStore_NonNilValue(t *testing.T) {
 	}
 }
 
+func TestMapLoadAndStore_NilValue(t *testing.T) {
+	m := NewMap()
+	m.LoadAndStore("foo", nil)
+	v, loaded := m.LoadAndStore("foo", nil)
+	if !loaded {
+		t.Error("nil value was expected")
+	}
+	if v != nil {
+		t.Errorf("value was not nil: %v", v)
+	}
+	v, loaded = m.Load("foo")
+	if !loaded {
+		t.Error("nil value was expected")
+	}
+	if v != nil {
+		t.Errorf("value was not nil: %v", v)
+	}
+}
+
+func TestMapLoadAndStore_NonNilValue(t *testing.T) {
+	type foo struct{}
+	m := NewMap()
+	v1 := &foo{}
+	v, loaded := m.LoadAndStore("foo", v1)
+	if loaded {
+		t.Error("no value was expected")
+	}
+	if v != v1 {
+		t.Errorf("value does not match: %v", v)
+	}
+	v2 := 2
+	v, loaded = m.LoadAndStore("foo", v2)
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != v1 {
+		t.Errorf("value does not match: %v", v)
+	}
+	v, loaded = m.Load("foo")
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != v2 {
+		t.Errorf("value does not match: %v", v)
+	}
+}
+
 func TestMapRange(t *testing.T) {
 	const numEntries = 1000
 	m := NewMap()
diff --git a/mapof_test.go b/mapof_test.go
index 7f9fd1f..86f2d9a 100644
--- a/mapof_test.go
+++ b/mapof_test.go
@@ -75,6 +75,52 @@ func TestMapOfLoadOrStore_NonNilValue(t *testing.T) {
 	}
 }
 
+func TestMapOfLoadAndStore_NilValue(t *testing.T) {
+	m := NewMapOf[*struct{}]()
+	m.LoadAndStore("foo", nil)
+	v, loaded := m.LoadAndStore("foo", nil)
+	if !loaded {
+		t.Error("nil value was expected")
+	}
+	if v != nil {
+		t.Errorf("value was not nil: %v", v)
+	}
+	v, loaded = m.Load("foo")
+	if !loaded {
+		t.Error("nil value was expected")
+	}
+	if v != nil {
+		t.Errorf("value was not nil: %v", v)
+	}
+}
+
+func TestMapOfLoadAndStore_NonNilValue(t *testing.T) {
+	m := NewMapOf[int]()
+	v1 := 1
+	v, loaded := m.LoadAndStore("foo", v1)
+	if loaded {
+		t.Error("no value was expected")
+	}
+	if v != v1 {
+		t.Errorf("value does not match: %v", v)
+	}
+	v2 := 2
+	v, loaded = m.LoadAndStore("foo", v2)
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != v1 {
+		t.Errorf("value does not match: %v", v)
+	}
+	v, loaded = m.Load("foo")
+	if !loaded {
+		t.Error("value was expected")
+	}
+	if v != v2 {
+		t.Errorf("value does not match: %v", v)
+	}
+}
+
 func TestMapOfRange(t *testing.T) {
 	const numEntries = 1000
 	m := NewMapOf[int]()