4 files changed, 79 insertions, 97 deletions
diff --git a/lib/rbtree/print_test.go b/lib/rbtree/print_test.go
deleted file mode 100644
index fe2d2bd..0000000
--- a/lib/rbtree/print_test.go
+++ /dev/null
@@ -1,45 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rbtree
-
-import (
-	"fmt"
-	"io"
-	"strings"
-)
-
-func (t *Tree[K, V]) ASCIIArt() string {
-	var out strings.Builder
-	t.root.asciiArt(&out, "", "", "")
-	return out.String()
-}
-
-func (node *Node[V]) String() string {
-	switch {
-	case node == nil:
-		return "nil"
-	case node.Color == Red:
-		return fmt.Sprintf("R(%v)", node.Value)
-	default:
-		return fmt.Sprintf("B(%v)", node.Value)
-	}
-}
-
-func (node *Node[V]) asciiArt(w io.Writer, u, m, l string) {
-	if node == nil {
-		fmt.Fprintf(w, "%snil\n", m)
-		return
-	}
-
-	node.Right.asciiArt(w, u+"     ", u+"  ,--", u+"  |  ")
-
-	if node.Color == Red {
-		fmt.Fprintf(w, "%s%v\n", m, node)
-	} else {
-		fmt.Fprintf(w, "%s%v\n", m, node)
-	}
-
-	node.Left.asciiArt(w, l+"  |  ", l+"  `--", l+"     ")
-}
diff --git a/lib/rbtree/rbtree.go b/lib/rbtree/rbtree.go
index 575a9ca..25f3b1c 100644
--- a/lib/rbtree/rbtree.go
+++ b/lib/rbtree/rbtree.go
@@ -6,6 +6,7 @@ package rbtree
 
 import (
 	"fmt"
+	"reflect"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
@@ -180,6 +181,47 @@ func (cur *Node[V]) prev() *Node[V] {
 	return parent
 }
 
+// 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)
+}
+
 func (t *Tree[K, V]) parentChild(node *Node[V]) **Node[V] {
 	switch {
 	case node.Parent == nil:
diff --git a/lib/rbtree/rbtree_test.go b/lib/rbtree/rbtree_test.go
index 999eb6f..e6007e1 100644
--- a/lib/rbtree/rbtree_test.go
+++ b/lib/rbtree/rbtree_test.go
@@ -5,7 +5,10 @@
 package rbtree
 
 import (
+	"fmt"
+	"io"
 	"sort"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -14,6 +17,40 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
 
+func (t *Tree[K, V]) ASCIIArt() string {
+	var out strings.Builder
+	t.root.asciiArt(&out, "", "", "")
+	return out.String()
+}
+
+func (node *Node[V]) String() string {
+	switch {
+	case node == nil:
+		return "nil"
+	case node.Color == Red:
+		return fmt.Sprintf("R(%v)", node.Value)
+	default:
+		return fmt.Sprintf("B(%v)", node.Value)
+	}
+}
+
+func (node *Node[V]) asciiArt(w io.Writer, u, m, l string) {
+	if node == nil {
+		fmt.Fprintf(w, "%snil\n", m)
+		return
+	}
+
+	node.Right.asciiArt(w, u+"     ", u+"  ,--", u+"  |  ")
+
+	if node.Color == Red {
+		fmt.Fprintf(w, "%s%v\n", m, node)
+	} else {
+		fmt.Fprintf(w, "%s%v\n", m, node)
+	}
+
+	node.Left.asciiArt(w, l+"  |  ", l+"  `--", l+"     ")
+}
+
 func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]struct{}, tree *Tree[util.NativeOrdered[K], V]) {
 	// 1. Every node is either red or black
 
diff --git a/lib/rbtree/rbtree_util.go b/lib/rbtree/rbtree_util.go
deleted file mode 100644
index cee5508..0000000
--- a/lib/rbtree/rbtree_util.go
+++ /dev/null
@@ -1,52 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rbtree
-
-import (
-	"reflect"
-
-	"git.lukeshu.com/btrfs-progs-ng/lib/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)
-}