btrfsvol: use rbtree

3 files changed, 48 insertions, 13 deletions

diff --git a/pkg/rbtree/rbtree.go b/pkg/rbtree/rbtree.go
index 13d9602..7927307 100644
--- a/pkg/rbtree/rbtree.go
+++ b/pkg/rbtree/rbtree.go
@@ -33,17 +33,24 @@ type Tree[K constraints.Ordered, V any] struct {
 	root  *Node[V]
 }
 
-func (t *Tree[K, V]) Walk(fn func(*Node[V])) {
-	t.root.walk(fn)
+func (t *Tree[K, V]) Walk(fn func(*Node[V]) error) error {
+	return t.root.walk(fn)
 }
 
-func (node *Node[V]) walk(fn func(*Node[V])) {
+func (node *Node[V]) walk(fn func(*Node[V]) error) error {
 	if node == nil {
-		return
+		return nil
+	}
+	if err := node.Left.walk(fn); err != nil {
+		return err
+	}
+	if err := fn(node); err != nil {
+		return err
+	}
+	if err := node.Right.walk(fn); err != nil {
+		return err
 	}
-	node.Left.walk(fn)
-	fn(node)
-	node.Right.walk(fn)
+	return nil
 }
 
 // Search the tree for a value that satisfied the given callbackk
diff --git a/pkg/rbtree/rbtree_test.go b/pkg/rbtree/rbtree_test.go
index 39307f4..9b8e02c 100644
--- a/pkg/rbtree/rbtree_test.go
+++ b/pkg/rbtree/rbtree_test.go
@@ -17,12 +17,13 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 	// 3. Every nil is black.
 
 	// 4. If a node is red, then both its children are black.
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		if node.getColor() == Red {
 			require.Equal(t, Black, node.Left.getColor())
 			require.Equal(t, Black, node.Right.getColor())
 		}
-	})
+		return nil
+	}))
 
 	// 5. For each node, all simple paths from the node to
 	//    descendent leaves contain the same number of black
@@ -39,7 +40,7 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 		walkCnt(node.Left, cnt, leafFn)
 		walkCnt(node.Right, cnt, leafFn)
 	}
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		var cnts []int
 		walkCnt(node, 0, func(cnt int) {
 			cnts = append(cnts, cnt)
@@ -50,7 +51,8 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 				break
 			}
 		}
-	})
+		return nil
+	}))
 
 	// expected contents
 	expectedOrder := make([]K, 0, len(expectedSet))
@@ -64,9 +66,10 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 		return expectedOrder[i] < expectedOrder[j]
 	})
 	actOrder := make([]K, 0, len(expectedSet))
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		actOrder = append(actOrder, tree.KeyFn(node.Value))
-	})
+		return nil
+	}))
 	require.Equal(t, expectedOrder, actOrder)
 }
 
diff --git a/pkg/rbtree/range.go b/pkg/rbtree/rbtree_util.go
index 9f67078..db7fd32 100644
--- a/pkg/rbtree/range.go
+++ b/pkg/rbtree/rbtree_util.go
@@ -1,6 +1,8 @@
 package rbtree
 
 import (
+	"reflect"
+
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
@@ -21,3 +23,26 @@ func (t *Tree[K, V]) SearchRange(fn func(V) int) []V {
 	}
 	return ret
 }
+
+func (t *Tree[K, V]) Equal(u *Tree[K, V]) bool {
+	if (t == nil) != (u == nil) {
+		return false
+	}
+	if t == nil {
+		return true
+	}
+
+	var tSlice []V
+	t.Walk(func(node *Node[V]) error {
+		tSlice = append(tSlice, node.Value)
+		return nil
+	})
+
+	var uSlice []V
+	u.Walk(func(node *Node[V]) error {
+		uSlice = append(uSlice, node.Value)
+		return nil
+	})
+
+	return reflect.DeepEqual(tSlice, uSlice)
+}