[tools] Add containers.Map.GetOrCreate[Locked]() helpers
Change-Id: Ia8e8b0a87cbc1a9e6c084fe646c3862bc046fcf5
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/146381
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>
diff --git a/tools/src/container/map.go b/tools/src/container/map.go
index 6b8acdc..a1d1c49 100644
--- a/tools/src/container/map.go
+++ b/tools/src/container/map.go
@@ -14,7 +14,10 @@
package container
-import "sort"
+import (
+ "sort"
+ "sync"
+)
// Map is a generic unordered map, which wrap's go's builtin 'map'.
// K is the map key, which must match the 'key' constraint.
@@ -60,3 +63,33 @@
}
return out
}
+
+// GetOrCreate returns the value of the map entry with the given key, creating
+// the map entry with create() if the entry did not exist.
+func (m Map[K, V]) GetOrCreate(key K, create func() V) V {
+ value, ok := m[key]
+ if !ok {
+ value = create()
+ m[key] = value
+ }
+ return value
+}
+
+// GetOrCreateLocked is similar to GetOrCreate, but performs lookup with a
+// read-lock on the provided mutex, and a write lock on create() and map
+// insertion.
+func (m Map[K, V]) GetOrCreateLocked(mutex *sync.RWMutex, key K, create func() V) V {
+ mutex.RLock()
+ value, ok := m[key]
+ mutex.RUnlock()
+ if !ok {
+ mutex.Lock()
+ defer mutex.Unlock()
+ value, ok = m[key]
+ if !ok {
+ value = create()
+ m[key] = value
+ }
+ }
+ return value
+}
diff --git a/tools/src/container/map_test.go b/tools/src/container/map_test.go
index 292b428..f2fde04 100644
--- a/tools/src/container/map_test.go
+++ b/tools/src/container/map_test.go
@@ -15,6 +15,7 @@
package container_test
import (
+ "sync"
"testing"
"dawn.googlesource.com/dawn/tools/src/container"
@@ -112,3 +113,44 @@
m.Add("b", 3)
expectEq(t, `m.Values()`, m.Values(), []int{2, 3, 1})
}
+
+func TestMapGetOrCreate(t *testing.T) {
+ m := container.NewMap[string, int]()
+
+ A := m.GetOrCreate("1", func() int { return 1 })
+ expectEq(t, "A", A, 1)
+
+ B := m.GetOrCreate("2", func() int { return 2 })
+ expectEq(t, "B", B, 2)
+
+ C := m.GetOrCreate("1", func() int { t.Error("should not be called"); return 0 })
+ expectEq(t, "C", C, 1)
+
+ D := m.GetOrCreate("2", func() int { t.Error("should not be called"); return 0 })
+ expectEq(t, "D", D, 2)
+}
+
+func TestMapGetOrCreateLocked(t *testing.T) {
+ m := container.NewMap[string, int]()
+
+ mtx := &sync.RWMutex{}
+ wg := sync.WaitGroup{}
+ wg.Add(100)
+ for i := 0; i < 100; i++ {
+ go func() {
+ defer wg.Done()
+ A := m.GetOrCreateLocked(mtx, "1", func() int { return 1 })
+ expectEq(t, "A", A, 1)
+
+ B := m.GetOrCreateLocked(mtx, "2", func() int { return 2 })
+ expectEq(t, "B", B, 2)
+
+ C := m.GetOrCreateLocked(mtx, "1", func() int { t.Error("should not be called"); return 0 })
+ expectEq(t, "C", C, 1)
+
+ D := m.GetOrCreateLocked(mtx, "2", func() int { t.Error("should not be called"); return 0 })
+ expectEq(t, "D", D, 2)
+ }()
+ }
+ wg.Wait()
+}
