add UpdateOrInsert

Author: guonaihong

cmap/cmap.go

@@ -96,6 +96,19 @@ func (c *CMap[K, V]) Delete(key K) {
 	item.rw.Unlock()
 }
 
+type UpdataOrInsertCb[K constraints.Ordered, V any] func(exits bool, old V) (newVal V)
+
+// 删除或者更新
+func (c *CMap[K, V]) UpdateOrInsert(k K, cb UpdataOrInsertCb[K, V]) {
+	item := c.findIndex(k)
+	item.rw.Lock()
+	old, ok := item.m.GetWithBool(k)
+	newVal := cb(ok, old)
+	item.m.Set(k, newVal)
+	item.rw.Unlock()
+
+}
+
 func (c *CMap[K, V]) Load(key K) (value V, ok bool) {
 	item := c.findIndex(key)
 	item.rw.RLock()

cmap/cmap_test.go

@@ -182,3 +182,41 @@ func Test_Values(t *testing.T) {
 	m2 := New[string, string]()
 	assert.Equal(t, len(m2.Keys()), 0)
 }
+
+func Test_UpdateOrInsert(t *testing.T) {
+	t.Run("Update", func(t *testing.T) {
+		m := New[string, string]()
+		m.Store("a", "1")
+		m.Store("b", "2")
+		m.Store("c", "3")
+		m.UpdateOrInsert("a", func(exist bool, old string) string {
+			if !exist {
+				t.Error("should exist")
+			}
+			if exist {
+				return "4"
+			}
+			return old
+		})
+		get, _ := m.Load("a")
+		if get != "4" {
+			t.Error("should be 4")
+		}
+	})
+
+	t.Run("Insert", func(t *testing.T) {
+		m := New[string, string]()
+		m.Store("a", "1")
+		m.UpdateOrInsert("b", func(exist bool, old string) string {
+			if !exist {
+				return "2"
+			}
+			return ""
+		})
+
+		get, _ := m.Load("b")
+		if get != "2" {
+			t.Error("should be 2")
+		}
+	})
+}

rwmap/rwmap.go

@@ -9,8 +9,7 @@ import (
 	"github.com/antlabs/gstl/mapex"
 )
 
-//type Pair[K comparable, V any] = mapex.Pair[K comparable, V any]
-
+// type Pair[K comparable, V any] = mapex.Pair[K comparable, V any]
 type Pair[K comparable, V any] struct {
 	Key K
 	Val V
@@ -47,7 +46,6 @@ func (r *RWMap[K, V]) Load(key K) (value V, ok bool) {
 	value, ok = r.m[key]
 	r.rw.RUnlock()
 	return
-
 }
 
 // 获取值，然后并删除
@@ -98,7 +96,6 @@ func (r *RWMap[K, V]) Iter() <-chan Pair[K, V] {
 		}
 		close(p)
 		r.rw.RUnlock()
-
 	}()
 	return p
 }
@@ -115,7 +112,6 @@ func (r *RWMap[K, V]) Store(key K, value V) {
 
 // keys
 func (r *RWMap[K, V]) Keys() (keys []K) {
-
 	r.rw.RLock()
 	if r.m == nil {
 		r.rw.RUnlock()
@@ -128,7 +124,6 @@ func (r *RWMap[K, V]) Keys() (keys []K) {
 
 // vals
 func (r *RWMap[K, V]) Values() (values []V) {
-
 	r.rw.RLock()
 	if r.m == nil {
 		r.rw.RUnlock()