From 3aa371a6d97314b7309802f3dc8a389bac7fdd63 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      | 24 ++++++++++++++++
 concurrent_map_test.go | 62 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/concurrent_map.go b/concurrent_map.go
index baccab0..44faf1a 100644
--- a/concurrent_map.go
+++ b/concurrent_map.go
@@ -75,6 +75,30 @@ 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.
+// Returning false for the second return value aborts the update.
+// 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(exist bool, valueInMap V) (V, bool)
+
+// 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.
+// if the callback return false for the second return value, the map will not be updated.
+func (m ConcurrentMap[K, V]) Update(key K, cb UpdateCb[V]) (res V) {
+	shard := m.GetShard(key)
+	shard.Lock()
+	v, ok := shard.items[key]
+	res, update := cb(ok, v)
+	if update {
+		shard.items[key] = res
+	}
+	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..a62c52b 100644
--- a/concurrent_map_test.go
+++ b/concurrent_map_test.go
@@ -479,6 +479,68 @@ func TestFnv32(t *testing.T) {
 
 }
 
+func TestUpdate(t *testing.T) {
+	m := New[Animal]()
+	lion := Animal{"lion"}
+
+	m.Set("safari", lion)
+	m.Update("safari", func(exists bool, valueInMap Animal) (Animal, bool) {
+		if !exists {
+			t.Error("Update recieved false exists flag for existing key")
+		}
+		valueInMap.name = "tiger"
+		return valueInMap, true
+	})
+	safari, ok := m.Get("safari")
+	if safari.name != "tiger" || !ok {
+		t.Error("Set, then Update failed")
+	}
+
+	m.Update("marine", func(exists bool, valueInMap Animal) (Animal, bool) {
+		if exists {
+			t.Error("Update recieved exists flag for empty key")
+		}
+		if valueInMap.name != "" {
+			t.Error("Update did not receive zero value for non existing key")
+		}
+		valueInMap.name = "whale"
+		return valueInMap, true
+	})
+	marineAnimals, ok := m.Get("marine")
+	if marineAnimals.name != "whale" || !ok {
+		t.Error("Update on non-existing key failed")
+	}
+
+	// return false to prevent updateing map
+	m.Set("safari", lion)
+	m.Update("safari", func(exists bool, valueInMap Animal) (Animal, bool) {
+		if !exists {
+			t.Error("Update recieved false exists flag for existing key")
+		}
+		valueInMap.name = "tiger"
+		return valueInMap, false
+	})
+	safari, ok = m.Get("safari")
+	if safari.name != "lion" || !ok {
+		t.Error("Set, then aborting Update failed")
+	}
+
+	m.Update("tundra", func(exists bool, valueInMap Animal) (Animal, bool) {
+		if exists {
+			t.Error("Update recieved exists flag for empty key")
+		}
+		if valueInMap.name != "" {
+			t.Error("Update did not receive zero value for non existing key")
+		}
+		valueInMap.name = "moose"
+		return valueInMap, false
+	})
+	_, ok = m.Get("tundra")
+	if ok {
+		t.Error("Update aborting on non-existing key failed")
+	}
+}
+
 func TestUpsert(t *testing.T) {
 	dolphin := Animal{"dolphin"}
 	whale := Animal{"whale"}