Skip to content

Commit

Permalink
add UpdateOrInsert
Browse files Browse the repository at this point in the history
  • Loading branch information
guonaihong committed Jul 30, 2024
1 parent 490120d commit 957b5c4
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 6 deletions.
13 changes: 13 additions & 0 deletions cmap/cmap.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
38 changes: 38 additions & 0 deletions cmap/cmap_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
})
}
7 changes: 1 addition & 6 deletions rwmap/rwmap.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

}

// 获取值,然后并删除
Expand Down Expand Up @@ -98,7 +96,6 @@ func (r *RWMap[K, V]) Iter() <-chan Pair[K, V] {
}
close(p)
r.rw.RUnlock()

}()
return p
}
Expand All @@ -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()
Expand All @@ -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()
Expand Down

0 comments on commit 957b5c4

Please sign in to comment.