feat(csync.Map): added GetOrSet

Carlos Alexandro Becker created 2 months ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/csync/maps.go      | 12 ++++++++++++
internal/csync/maps_test.go | 10 ++++++++++
2 files changed, 22 insertions(+)

Detailed changes

internal/csync/maps.go

@@ -56,6 +56,18 @@ func (m *Map[K, V]) Len() int {
 	return len(m.inner)
 }
 
+// GetOrSet gets and returns the key if it exists, otherwise, it executes the
+// given function, set its return value for the given key, and returns it.
+func (m *Map[K, V]) GetOrSet(key K, fn func() V) V {
+	got, ok := m.Get(key)
+	if ok {
+		return got
+	}
+	value := fn()
+	m.Set(key, value)
+	return value
+}
+
 // Take gets an item and then deletes it.
 func (m *Map[K, V]) Take(key K) (V, bool) {
 	m.mu.Lock()

internal/csync/maps_test.go

@@ -54,6 +54,16 @@ func TestMap_Set(t *testing.T) {
 	require.Equal(t, 1, m.Len())
 }
 
+func TestMap_GetOrSet(t *testing.T) {
+	t.Parallel()
+
+	m := NewMap[string, int]()
+
+	require.Equal(t, 42, m.GetOrSet("key1", func() int { return 42 }))
+	require.Equal(t, 42, m.GetOrSet("key1", func() int { return 99999 }))
+	require.Equal(t, 1, m.Len())
+}
+
 func TestMap_Get(t *testing.T) {
 	t.Parallel()