From 27d2f3a0efe6de94c7720907557e640e8a2f1428 Mon Sep 17 00:00:00 2001 From: Luke Shumaker Date: Thu, 30 Jun 2022 03:14:05 -0600 Subject: btrfsvol: use rbtree --- pkg/rbtree/range.go | 23 ----------------------- pkg/rbtree/rbtree.go | 21 ++++++++++++++------- pkg/rbtree/rbtree_test.go | 15 +++++++++------ pkg/rbtree/rbtree_util.go | 48 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 36 deletions(-) delete mode 100644 pkg/rbtree/range.go create mode 100644 pkg/rbtree/rbtree_util.go (limited to 'pkg/rbtree') diff --git a/pkg/rbtree/range.go b/pkg/rbtree/range.go deleted file mode 100644 index 9f67078..0000000 --- a/pkg/rbtree/range.go +++ /dev/null @@ -1,23 +0,0 @@ -package rbtree - -import ( - "lukeshu.com/btrfs-tools/pkg/util" -) - -// SearchRange is like Search, but returns all nodes that match the -// function; assuming that they are contiguous. -func (t *Tree[K, V]) SearchRange(fn func(V) int) []V { - middle := t.Search(fn) - if middle == nil { - return nil - } - ret := []V{middle.Value} - for node := t.Prev(middle); node != nil && fn(node.Value) == 0; node = t.Prev(node) { - ret = append(ret, node.Value) - } - util.ReverseSlice(ret) - for node := t.Next(middle); node != nil && fn(node.Value) == 0; node = t.Next(node) { - ret = append(ret, node.Value) - } - return ret -} 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/rbtree_util.go b/pkg/rbtree/rbtree_util.go new file mode 100644 index 0000000..db7fd32 --- /dev/null +++ b/pkg/rbtree/rbtree_util.go @@ -0,0 +1,48 @@ +package rbtree + +import ( + "reflect" + + "lukeshu.com/btrfs-tools/pkg/util" +) + +// SearchRange is like Search, but returns all nodes that match the +// function; assuming that they are contiguous. +func (t *Tree[K, V]) SearchRange(fn func(V) int) []V { + middle := t.Search(fn) + if middle == nil { + return nil + } + ret := []V{middle.Value} + for node := t.Prev(middle); node != nil && fn(node.Value) == 0; node = t.Prev(node) { + ret = append(ret, node.Value) + } + util.ReverseSlice(ret) + for node := t.Next(middle); node != nil && fn(node.Value) == 0; node = t.Next(node) { + ret = append(ret, node.Value) + } + 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) +} -- cgit v1.2.3-2-g168b