containers: rbtree.go: Track size of the tree

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) {