feat: add foreach and clear to lazymap

Adds foreach, remove, getinverse, removeinverse and clear methods to LazyMap, BiMap and LazyBiMap.

Signed-off-by: Utku Ozdemir <utku.ozdemir@siderolabs.com>

Author: utkuozdemir

containers/lazymap.go

@@ -35,6 +35,33 @@ func (m *LazyBiMap[K, V]) Get(key K) (V, bool) {
 	return val, ok
 }
 
+// GetInverse returns the key for the given value.
+func (m *LazyBiMap[K, V]) GetInverse(value V) (K, bool) {
+	key, ok := m.biMap.GetInverse(value)
+
+	return key, ok
+}
+
+// Remove removes the value for the given key.
+func (m *LazyBiMap[K, V]) Remove(key K) {
+	m.biMap.Remove(key)
+}
+
+// RemoveInverse removes the key for the given value.
+func (m *LazyBiMap[K, V]) RemoveInverse(value V) {
+	m.biMap.RemoveInverse(value)
+}
+
+// Clear removes all values.
+func (m *LazyBiMap[K, V]) Clear() {
+	m.biMap.Clear()
+}
+
+// ForEach calls the given function for each key-value pair.
+func (m *LazyBiMap[K, V]) ForEach(f func(K, V)) {
+	m.biMap.ForEach(f)
+}
+
 // BiMap (or “bidirectional map”) is a special kind of map that maintains
 // an inverse view of the map while ensuring that no duplicate values are present
 // and a value can always be used safely to get the key back.
@@ -50,6 +77,13 @@ func (m *BiMap[K, V]) Get(key K) (V, bool) {
 	return val, ok
 }
 
+// GetInverse returns the key for the given value.
+func (m *BiMap[K, V]) GetInverse(val V) (K, bool) {
+	key, ok := m.v2k[val]
+
+	return key, ok
+}
+
 // Set sets the value for the given key.
 func (m *BiMap[K, V]) Set(key K, val V) {
 	if m.k2v == nil {
@@ -66,19 +100,62 @@ func (m *BiMap[K, V]) Set(key K, val V) {
 	m.v2k[val] = key
 }
 
+// Remove deletes the value for the given key.
+func (m *BiMap[K, V]) Remove(key K) {
+	if m.k2v == nil {
+		return
+	}
+
+	val, ok := m.k2v[key]
+	if !ok {
+		return
+	}
+
+	delete(m.k2v, key)
+	delete(m.v2k, val)
+}
+
+// RemoveInverse deletes the key for the given value.
+func (m *BiMap[K, V]) RemoveInverse(val V) {
+	if m.v2k == nil {
+		return
+	}
+
+	key, ok := m.v2k[val]
+	if !ok {
+		return
+	}
+
+	delete(m.v2k, val)
+	delete(m.k2v, key)
+}
+
+// Clear removes all key-value pairs.
+func (m *BiMap[K, V]) Clear() {
+	m.k2v = nil
+	m.v2k = nil
+}
+
+// ForEach calls the given function for each key-value pair.
+func (m *BiMap[K, V]) ForEach(f func(K, V)) {
+	for k, v := range m.k2v {
+		f(k, v)
+	}
+}
+
 // LazyMap is like usual map but creates values on demand.
 type LazyMap[K comparable, V comparable] struct {
 	Creator func(K) (V, error)
-	biMap   map[K]V
+	dataMap map[K]V
 }
 
 // GetOrCreate returns the value for the given key. It creates it using Creator if it doesn't exist.
 func (m *LazyMap[K, V]) GetOrCreate(key K) (V, error) {
-	if m.biMap == nil {
-		m.biMap = map[K]V{}
+	if m.dataMap == nil {
+		m.dataMap = map[K]V{}
 	}
 
-	val, ok := m.biMap[key]
+	val, ok := m.dataMap[key]
 	if ok {
 		return val, nil
 	}
@@ -88,23 +165,35 @@ func (m *LazyMap[K, V]) GetOrCreate(key K) (V, error) {
 		return *new(V), err //nolint:gocritic
 	}
 
-	m.biMap[key] = val
+	m.dataMap[key] = val
 
 	return val, nil
 }
 
 // Get returns the value for the given key.
 func (m *LazyMap[K, V]) Get(key K) (V, bool) {
-	val, ok := m.biMap[key]
+	val, ok := m.dataMap[key]
 
 	return val, ok
 }
 
 // Remove deletes the value for the given key.
 func (m *LazyMap[K, V]) Remove(key K) {
-	if m.biMap == nil {
+	if m.dataMap == nil {
 		return
 	}
 
-	delete(m.biMap, key)
+	delete(m.dataMap, key)
+}
+
+// Clear removes all key-value pairs.
+func (m *LazyMap[K, V]) Clear() {
+	m.dataMap = nil
+}
+
+// ForEach calls the given function for each key-value pair.
+func (m *LazyMap[K, V]) ForEach(f func(K, V)) {
+	for k, v := range m.dataMap {
+		f(k, v)
+	}
 }

containers/lazymap_test.go

@@ -2,12 +2,14 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+//nolint:dupl
 package containers_test
 
 import (
 	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/siderolabs/gen/containers"
@@ -32,9 +34,12 @@ func TestLazyBiMap(t *testing.T) {
 		},
 	}
 
-	t.Run("should return nothing if key doesnt exist", func(t *testing.T) {
+	t.Run("should return nothing if key or value doesnt exist", func(t *testing.T) {
 		_, ok := m.Get(0)
 		require.False(t, ok)
+
+		_, ok = m.GetInverse(0)
+		require.False(t, ok)
 	})
 
 	t.Run("should create value on demand", func(t *testing.T) {
@@ -49,6 +54,13 @@ func TestLazyBiMap(t *testing.T) {
 		require.Equal(t, 100, create)
 	})
 
+	t.Run("should return existing key by value", func(t *testing.T) {
+		key, ok := m.GetInverse(100)
+		require.True(t, ok)
+
+		assert.Equal(t, 1, key)
+	})
+
 	t.Run("should remove old key if has new key", func(t *testing.T) {
 		create, err := m.GetOrCreate(11)
 		require.NoError(t, err)
@@ -62,6 +74,67 @@ func TestLazyBiMap(t *testing.T) {
 		_, err := m.GetOrCreate(-1)
 		require.Error(t, err)
 	})
+
+	t.Run("should remove key", func(t *testing.T) {
+		_, err := m.GetOrCreate(12)
+		require.NoError(t, err)
+
+		m.Remove(12)
+		_, ok := m.Get(12)
+		require.False(t, ok)
+	})
+
+	t.Run("should remove key by value", func(t *testing.T) {
+		_, err := m.GetOrCreate(13)
+		require.NoError(t, err)
+
+		m.RemoveInverse(300)
+		_, ok := m.Get(13)
+		require.False(t, ok)
+	})
+
+	t.Run("should remove all entries", func(t *testing.T) {
+		_, err := m.GetOrCreate(14)
+		require.NoError(t, err)
+
+		_, err = m.GetOrCreate(15)
+		require.NoError(t, err)
+
+		m.Clear()
+
+		_, ok := m.Get(13)
+		require.False(t, ok)
+
+		_, ok = m.Get(14)
+		require.False(t, ok)
+	})
+
+	t.Run("should iterate over entries", func(t *testing.T) {
+		var keys []int
+		var values []int
+
+		m.ForEach(func(k int, v int) {
+			keys = append(keys, k)
+			values = append(values, v)
+		})
+
+		assert.Empty(t, keys)
+		assert.Empty(t, values)
+
+		_, err := m.GetOrCreate(1)
+		require.NoError(t, err)
+
+		_, err = m.GetOrCreate(2)
+		require.NoError(t, err)
+
+		m.ForEach(func(k int, v int) {
+			keys = append(keys, k)
+			values = append(values, v)
+		})
+
+		assert.Equal(t, []int{1, 2}, keys)
+		assert.Equal(t, []int{100, 200}, values)
+	})
 }
 
 func TestLazyMap(t *testing.T) {
@@ -114,4 +187,47 @@ func TestLazyMap(t *testing.T) {
 		_, ok := m.Get(1)
 		require.False(t, ok)
 	})
+
+	t.Run("should remove all entries", func(t *testing.T) {
+		_, err := m.GetOrCreate(2)
+		require.NoError(t, err)
+
+		_, err = m.GetOrCreate(3)
+		require.NoError(t, err)
+
+		m.Clear()
+
+		_, ok := m.Get(1)
+		require.False(t, ok)
+
+		_, ok = m.Get(2)
+		require.False(t, ok)
+	})
+
+	t.Run("should iterate over entries", func(t *testing.T) {
+		var keys []int
+		var values []int
+
+		m.ForEach(func(k int, v int) {
+			keys = append(keys, k)
+			values = append(values, v)
+		})
+
+		assert.Empty(t, keys)
+		assert.Empty(t, values)
+
+		_, err := m.GetOrCreate(4)
+		require.NoError(t, err)
+
+		_, err = m.GetOrCreate(5)
+		require.NoError(t, err)
+
+		m.ForEach(func(k int, v int) {
+			keys = append(keys, k)
+			values = append(values, v)
+		})
+
+		assert.Equal(t, []int{4, 5}, keys)
+		assert.Equal(t, []int{400, 500}, values)
+	})
 }