1 files changed, 24 insertions, 52 deletions

diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go
index 52308c9..d104274 100644
--- a/lib/containers/sortedmap.go
+++ b/lib/containers/sortedmap.go
@@ -1,40 +1,32 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 package containers
 
-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]]
-}
-
-func (m *SortedMap[K, V]) init() {
-	if m.inner.KeyFn == nil {
-		m.inner.KeyFn = m.keyFn
-	}
+func (a orderedKV[K, V]) Compare(b orderedKV[K, V]) int {
+	return a.K.Compare(b.K)
 }
 
-func (m *SortedMap[K, V]) keyFn(kv OrderedKV[K, V]) K {
-	return kv.K
+type SortedMap[K Ordered[K], V any] struct {
+	inner RBTree[orderedKV[K, V]]
 }
 
 func (m *SortedMap[K, V]) Delete(key K) {
-	m.init()
-	m.inner.Delete(key)
+	m.inner.Delete(m.inner.Search(func(kv orderedKV[K, V]) int {
+		return key.Compare(kv.K)
+	}))
 }
 
 func (m *SortedMap[K, V]) Load(key K) (value V, ok bool) {
-	m.init()
-	node := m.inner.Lookup(key)
+	node := m.inner.Search(func(kv orderedKV[K, V]) int {
+		return key.Compare(kv.K)
+	})
 	if node == nil {
 		var zero V
 		return zero, false
@@ -42,41 +34,27 @@ func (m *SortedMap[K, V]) Load(key K) (value V, ok bool) {
 	return node.Value.V, true
 }
 
-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 {
-		if f(node.Value.K, node.Value.V) {
-			return nil
-		} else {
-			return errStop
-		}
+func (m *SortedMap[K, V]) Store(key K, value V) {
+	m.inner.Insert(orderedKV[K, V]{
+		K: key,
+		V: value,
 	})
 }
 
-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 {
-		return rangeFn(kv.K, kv.V)
+func (m *SortedMap[K, V]) Range(fn func(key K, value V) bool) {
+	m.inner.Range(func(node *RBNode[orderedKV[K, V]]) bool {
+		return fn(node.Value.K, node.Value.V)
 	})
-	for _, kv := range kvs {
-		if !handleFn(kv.K, kv.V) {
-			break
-		}
-	}
 }
 
-func (m *SortedMap[K, V]) Store(key K, value V) {
-	m.init()
-	m.inner.Insert(OrderedKV[K, V]{
-		K: key,
-		V: value,
-	})
+func (m *SortedMap[K, V]) Subrange(rangeFn func(K, V) int, handleFn func(K, V) bool) {
+	m.inner.Subrange(
+		func(kv orderedKV[K, V]) int { return rangeFn(kv.K, kv.V) },
+		func(node *RBNode[orderedKV[K, V]]) bool { return handleFn(node.Value.K, node.Value.V) })
 }
 
 func (m *SortedMap[K, V]) Search(fn func(K, V) int) (K, V, bool) {
-	node := m.inner.Search(func(kv OrderedKV[K, V]) int {
+	node := m.inner.Search(func(kv orderedKV[K, V]) int {
 		return fn(kv.K, kv.V)
 	})
 	if node == nil {
@@ -87,12 +65,6 @@ func (m *SortedMap[K, V]) Search(fn func(K, V) int) (K, V, bool) {
 	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)
-	})
-}
-
 func (m *SortedMap[K, V]) Len() int {
 	return m.inner.Len()
 }