Skip to content

Commit

Permalink
feat: add foreach and clear to lazymap
Browse files Browse the repository at this point in the history 
Adds foreach, remove, getinverse, removeinverse and clear methods to LazyMap, BiMap and LazyBiMap.

Signed-off-by: Utku Ozdemir <utku.ozdemir@siderolabs.com>
  • Loading branch information
@utkuozdemir
utkuozdemir committed Feb 6, 2023
1 parent 214c1ef commit 476dfea
Show file tree
Hide file tree
Showing 2 changed files with 214 additions and 9 deletions.
105 changes: 97 additions & 8 deletions containers/lazymap.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 {
Expand All @@ -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
}
Expand All @@ -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)
}
}
118 changes: 117 additions & 1 deletion containers/lazymap_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {
Expand All @@ -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)
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
})
}

0 comments on commit 476dfea

Please sign in to comment.