From 476dfeae70882e1ca6e5cfed3d6e12dc36841a26 Mon Sep 17 00:00:00 2001 From: Utku Ozdemir Date: Mon, 6 Feb 2023 11:20:14 +0100 Subject: [PATCH] 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 --- containers/lazymap.go | 105 ++++++++++++++++++++++++++++++--- containers/lazymap_test.go | 118 ++++++++++++++++++++++++++++++++++++- 2 files changed, 214 insertions(+), 9 deletions(-) diff --git a/containers/lazymap.go b/containers/lazymap.go index cb7d6ea..ab023fa 100644 --- a/containers/lazymap.go +++ b/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) + } } diff --git a/containers/lazymap_test.go b/containers/lazymap_test.go index 777105b..e1992c7 100644 --- a/containers/lazymap_test.go +++ b/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) + }) }