From ba8662ac09e38aba49ac33da82dd4f564faa3dbf Mon Sep 17 00:00:00 2001
From: Tomer Brisker <tbrisker@gmail.com>
Date: Sun, 2 Apr 2023 15:52:50 +0300
Subject: [PATCH] Implement Update method

The update method allows atomically mutating a value in the map.
If the key does not exist, the zero value for the value type is mutated
and saved.
---
 concurrent_map.go      | 19 +++++++++++++++++++
 concurrent_map_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/concurrent_map.go b/concurrent_map.go
index baccab0..16436b3 100644
--- a/concurrent_map.go
+++ b/concurrent_map.go
@@ -75,6 +75,25 @@ func (m ConcurrentMap[K, V]) Set(key K, value V) {
 	shard.Unlock()
 }
 
+// Callback to update an element in the map.
+// If the element doesn't exist in the map, the parameter will receive the zero value for the value type.
+// The returned value will be stored in the map replacing the existing value.
+// It is called while lock is held, therefore it MUST NOT
+// try to access other keys in same map, as it can lead to deadlock since
+// Go sync.RWLock is not reentrant
+type UpdateCb[V any] func(valueInMap V) V
+
+// Update an existing element using UpdateCb, assuming the key exists.
+// If it does not, the zero value for the value type is passed to the callback.
+func (m ConcurrentMap[K, V]) Update(key K, cb UpdateCb[V]) (res V) {
+	shard := m.GetShard(key)
+	shard.Lock()
+	v := shard.items[key]
+	shard.items[key] = cb(v)
+	shard.Unlock()
+	return res
+}
+
 // Callback to return new element to be inserted into the map
 // It is called while lock is held, therefore it MUST NOT
 // try to access other keys in same map, as it can lead to deadlock since
diff --git a/concurrent_map_test.go b/concurrent_map_test.go
index 4e5d52e..0de5ca2 100644
--- a/concurrent_map_test.go
+++ b/concurrent_map_test.go
@@ -479,6 +479,34 @@ func TestFnv32(t *testing.T) {
 
 }
 
+func TestUpdate(t *testing.T) {
+	m := New[Animal]()
+	lion := Animal{"lion"}
+
+	m.Set("safari", lion)
+	m.Update("safari", func(valueInMap Animal) Animal {
+		valueInMap.name = "tiger"
+		return valueInMap
+	})
+	safari, ok := m.Get("safari")
+	if safari.name != "tiger" || !ok {
+		t.Error("Set, then Update failed")
+	}
+
+	m.Update("marine", func(valueInMap Animal) Animal {
+		if valueInMap.name != "" {
+			t.Error("Update did not receive zero value for non existing key")
+		}
+		valueInMap.name = "whale"
+		return valueInMap
+	})
+	marineAnimals, ok := m.Get("marine")
+	if marineAnimals.name != "whale" || !ok {
+		t.Error("Update on non-existing key failed")
+	}
+
+}
+
 func TestUpsert(t *testing.T) {
 	dolphin := Animal{"dolphin"}
 	whale := Animal{"whale"}