summaryrefslogtreecommitdiff
path: root/lib/containers
diff options
context:
space:
mode:
Diffstat (limited to 'lib/containers')
-rw-r--r--lib/containers/rbtree.go7
-rw-r--r--lib/containers/rbtree_test.go1
2 files changed, 8 insertions, 0 deletions
diff --git a/lib/containers/rbtree.go b/lib/containers/rbtree.go
index ab79a26..e2d8f8a 100644
--- a/lib/containers/rbtree.go
+++ b/lib/containers/rbtree.go
@@ -37,6 +37,11 @@ type RBTree[K Ordered[K], V any] struct {
KeyFn func(V) K
AttrFn func(*RBNode[V])
root *RBNode[V]
+ len int
+}
+
+func (t *RBTree[K, V]) Len() int {
+ return t.len
}
func (t *RBTree[K, V]) Walk(fn func(*RBNode[V]) error) error {
@@ -324,6 +329,7 @@ func (t *RBTree[K, V]) Insert(val V) {
exact.Value = val
return
}
+ t.len++
node := &RBNode[V]{
Color: Red,
@@ -394,6 +400,7 @@ func (t *RBTree[K, V]) Delete(key K) {
if nodeToDelete == nil {
return
}
+ t.len--
// This is closely based on the algorithm presented in CLRS
// 3e.
diff --git a/lib/containers/rbtree_test.go b/lib/containers/rbtree_test.go
index 9841d26..a487c52 100644
--- a/lib/containers/rbtree_test.go
+++ b/lib/containers/rbtree_test.go
@@ -112,6 +112,7 @@ func checkRBTree[K constraints.Ordered, V any](t *testing.T, expectedSet Set[K],
return nil
}))
require.Equal(t, expectedOrder, actOrder)
+ require.Equal(t, tree.Len(), len(expectedSet))
}
func FuzzRBTree(f *testing.F) {