From b6b42fb33e84374180732af2ffcb16a513d45bea Mon Sep 17 00:00:00 2001
From: Hiram Chirino <hiram@hiramchirino.com>
Date: Sun, 5 Nov 2023 09:17:33 -0500
Subject: [PATCH] feat: add SetIf and DeleteIf methods.

These methods allow you to lazy create/initialize values and decide if
the values should be set or deleted while the map shard is locked.

Signed-off-by: Hiram Chirino <hiram@hiramchirino.com>
---
 concurrent_swiss_map.go      | 24 +++++++++++++++
 concurrent_swiss_map_test.go | 60 +++++++++++++++++++++++++++++++++++-
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/concurrent_swiss_map.go b/concurrent_swiss_map.go
index ac87e23..6bf7553 100644
--- a/concurrent_swiss_map.go
+++ b/concurrent_swiss_map.go
@@ -79,6 +79,18 @@ func (m *CsMap[K, V]) Delete(key K) bool {
 	return shard.items.DeleteWithHash(key, hashShardPair.hash)
 }
 
+func (m *CsMap[K, V]) DeleteIf(key K, condition func(value V) bool) bool {
+	hashShardPair := m.getShard(key)
+	shard := hashShardPair.shard
+	shard.Lock()
+	defer shard.Unlock()
+	value, ok := shard.items.GetWithHash(key, hashShardPair.hash)
+	if ok && condition(value) {
+		return shard.items.DeleteWithHash(key, hashShardPair.hash)
+	}
+	return false
+}
+
 func (m *CsMap[K, V]) Load(key K) (V, bool) {
 	hashShardPair := m.getShard(key)
 	shard := hashShardPair.shard
@@ -127,6 +139,18 @@ func (m *CsMap[K, V]) SetIfAbsent(key K, value V) {
 	}
 }
 
+func (m *CsMap[K, V]) SetIf(key K, conditionFn func(previousVale V, previousFound bool) (value V, set bool)) {
+	hashShardPair := m.getShard(key)
+	shard := hashShardPair.shard
+	shard.Lock()
+	defer shard.Unlock()
+	value, found := shard.items.GetWithHash(key, hashShardPair.hash)
+	value, ok := conditionFn(value, found)
+	if ok {
+		shard.items.PutWithHash(key, value, hashShardPair.hash)
+	}
+}
+
 func (m *CsMap[K, V]) SetIfPresent(key K, value V) {
 	hashShardPair := m.getShard(key)
 	shard := hashShardPair.shard
diff --git a/concurrent_swiss_map_test.go b/concurrent_swiss_map_test.go
index 21bd331..3b6f66c 100644
--- a/concurrent_swiss_map_test.go
+++ b/concurrent_swiss_map_test.go
@@ -52,7 +52,6 @@ func TestSetIfAbsent(t *testing.T) {
 		t.Fatal("1 should be exist")
 	}
 }
-
 func TestSetIfPresent(t *testing.T) {
 	myMap := csmap.Create[int, string]()
 	myMap.SetIfPresent(1, "test")
@@ -68,6 +67,65 @@ func TestSetIfPresent(t *testing.T) {
 	}
 }
 
+func TestSetIf(t *testing.T) {
+	myMap := csmap.Create[int, string]()
+	myMap.SetIf(1, func(previousVale string, previousFound bool) (value string, set bool) {
+		// operate like  a SetIfAbsent...
+		if !previousFound {
+			return "value a", true
+		}
+		return "", false
+	})
+	value, _ := myMap.Load(1)
+	if value != "value a" {
+		t.Fatal("value should value a")
+	}
+
+	myMap.SetIf(1, func(previousVale string, previousFound bool) (value string, set bool) {
+		// operate like  a SetIfAbsent...
+		if !previousFound {
+			return "bad", true
+		}
+		return "", false
+	})
+	value, _ = myMap.Load(1)
+	if value != "value a" {
+		t.Fatal("value should value a")
+	}
+}
+
+func TestDeleteIf(t *testing.T) {
+	myMap := csmap.Create[int, string]()
+	myMap.Store(1, "value b")
+	ok1 := myMap.DeleteIf(20, func(value string) bool {
+		t.Fatal("condition function should not have been called")
+		return false
+	})
+	if ok1 {
+		t.Fatal("ok1 should be false")
+	}
+
+	ok2 := myMap.DeleteIf(1, func(value string) bool {
+		if value != "value b" {
+			t.Fatal("condition function arg should be tests")
+		}
+		return false // don't delete
+	})
+	if ok2 {
+		t.Fatal("ok1 should be false")
+	}
+
+	ok3 := myMap.DeleteIf(1, func(value string) bool {
+		if value != "value b" {
+			t.Fatal("condition function arg should be tests")
+		}
+		return true // delete the entry
+	})
+	if !ok3 {
+		t.Fatal("ok2 should be true")
+	}
+}
+
 func TestCount(t *testing.T) {
 	myMap := csmap.Create[int, string]()
 	myMap.SetIfAbsent(1, "test")