summaryrefslogtreecommitdiff
path: root/lib/containers/sortedmap.go
diff options
context:
space:
mode:
Diffstat (limited to 'lib/containers/sortedmap.go')
-rw-r--r--lib/containers/sortedmap.go30
1 files changed, 24 insertions, 6 deletions
diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go
index b759fa3..d6d2f9e 100644
--- a/lib/containers/sortedmap.go
+++ b/lib/containers/sortedmap.go
@@ -8,13 +8,13 @@ import (
"errors"
)
-type orderedKV[K Ordered[K], V any] struct {
+type OrderedKV[K Ordered[K], V any] struct {
K K
V V
}
type SortedMap[K Ordered[K], V any] struct {
- inner RBTree[K, orderedKV[K, V]]
+ inner RBTree[K, OrderedKV[K, V]]
}
func (m *SortedMap[K, V]) init() {
@@ -23,7 +23,7 @@ func (m *SortedMap[K, V]) init() {
}
}
-func (m *SortedMap[K, V]) keyFn(kv orderedKV[K, V]) K {
+func (m *SortedMap[K, V]) keyFn(kv OrderedKV[K, V]) K {
return kv.K
}
@@ -46,7 +46,7 @@ var errStop = errors.New("stop")
func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) {
m.init()
- _ = m.inner.Walk(func(node *RBNode[orderedKV[K, V]]) error {
+ _ = m.inner.Walk(func(node *RBNode[OrderedKV[K, V]]) error {
if f(node.Value.K, node.Value.V) {
return nil
} else {
@@ -57,7 +57,7 @@ func (m *SortedMap[K, V]) Range(f func(key K, value V) bool) {
func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) bool) {
m.init()
- kvs := m.inner.SearchRange(func(kv orderedKV[K, V]) int {
+ kvs := m.inner.SearchRange(func(kv OrderedKV[K, V]) int {
return rangeFn(kv.K, kv.V)
})
for _, kv := range kvs {
@@ -69,8 +69,26 @@ func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) b
func (m *SortedMap[K, V]) Store(key K, value V) {
m.init()
- m.inner.Insert(orderedKV[K, V]{
+ m.inner.Insert(OrderedKV[K, V]{
K: key,
V: value,
})
}
+
+func (m *SortedMap[K, V]) Search(fn func(K, V) int) (K, V, bool) {
+ node := m.inner.Search(func(kv OrderedKV[K, V]) int {
+ return fn(kv.K, kv.V)
+ })
+ if node == nil {
+ var zeroK K
+ var zeroV V
+ return zeroK, zeroV, false
+ }
+ return node.Value.K, node.Value.V, true
+}
+
+func (m *SortedMap[K, V]) SearchAll(fn func(K, V) int) []OrderedKV[K, V] {
+ return m.inner.SearchRange(func(kv OrderedKV[K, V]) int {
+ return fn(kv.K, kv.V)
+ })
+}