Merge branch 'lukeshu/rebuild-nodes-take3'

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)
+	})
+}