modelRT/util/typed_map.go

85 lines
2.3 KiB
Go

// Package util provide some utility functions
package util
import (
"iter"
"sync"
)
// TypedMap define a type-safe generic wrapper around sync.Map.
// It keeps the concurrency guarantees of sync.Map while removing the
// per-call-site type assertions previously required for every Load.
type TypedMap[K comparable, V any] struct {
m sync.Map
}
// Load define func of return the value stored for key, and whether it was found.
func (t *TypedMap[K, V]) Load(key K) (V, bool) {
value, ok := t.m.Load(key)
if !ok {
var zero V
return zero, false
}
// safe: only values of type V are ever stored through this wrapper
return value.(V), true
}
// Store define func of set the value for key.
func (t *TypedMap[K, V]) Store(key K, value V) {
t.m.Store(key, value)
}
// LoadOrStore define func of return the existing value for key if present.
// Otherwise it stores and returns the given value. loaded reports whether the
// value was already present.
func (t *TypedMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
v, loaded := t.m.LoadOrStore(key, value)
return v.(V), loaded
}
// Swap define func of store value for key and return the previous value (if any).
// loaded reports whether a value was already present for key.
func (t *TypedMap[K, V]) Swap(key K, value V) (previous V, loaded bool) {
prev, loaded := t.m.Swap(key, value)
if !loaded {
var zero V
return zero, false
}
return prev.(V), true
}
// Delete define func of remove the value for key.
func (t *TypedMap[K, V]) Delete(key K) {
t.m.Delete(key)
}
// Range define func of iterate over all key/value pairs.
// Iteration stops early if f returns false.
func (t *TypedMap[K, V]) Range(f func(key K, value V) bool) {
t.m.Range(func(key, value any) bool {
return f(key.(K), value.(V))
})
}
// Len define func of return the number of key/value pairs currently stored.
// It walks the map, so the cost is O(n).
func (t *TypedMap[K, V]) Len() int {
count := 0
t.m.Range(func(_, _ any) bool {
count++
return true
})
return count
}
// All define func of return an iterator over all key/value pairs,
// usable directly with range-over-func. Iteration stops early if the
// consumer breaks out of the loop.
func (t *TypedMap[K, V]) All() iter.Seq2[K, V] {
return func(yield func(K, V) bool) {
t.m.Range(func(key, value any) bool {
return yield(key.(K), value.(V))
})
}
}