From 74f15aa6a5c04bab0615262d005eace4fc2baedb Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 30 Jun 2022 02:46:33 -0600
Subject: add rbtree.SearchRange

---
 pkg/rbtree/range.go | 23 +++++++++++++++++++++++
 pkg/util/generic.go |  7 +++++++
 2 files changed, 30 insertions(+)
 create mode 100644 pkg/rbtree/range.go

(limited to 'pkg')

diff --git a/pkg/rbtree/range.go b/pkg/rbtree/range.go
new file mode 100644
index 0000000..d7325ee
--- /dev/null
+++ b/pkg/rbtree/range.go
@@ -0,0 +1,23 @@
+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) []*Node[V] {
+	middle := t.Search(fn)
+	if middle == nil {
+		return nil
+	}
+	ret := []*Node[V]{middle}
+	for node := t.Prev(middle); node != nil && fn(node.Value) == 0; node = t.Prev(node) {
+		ret = append(ret, node)
+	}
+	util.ReverseSlice(ret)
+	for node := t.Next(middle); node != nil && fn(node.Value) == 0; node = t.Next(node) {
+		ret = append(ret, node)
+	}
+	return ret
+}
diff --git a/pkg/util/generic.go b/pkg/util/generic.go
index 520d94e..6474ea5 100644
--- a/pkg/util/generic.go
+++ b/pkg/util/generic.go
@@ -35,6 +35,13 @@ func RemoveAllFromSliceFunc[T any](haystack []T, f func(T) bool) []T {
 	return haystack
 }
 
+func ReverseSlice[T any](slice []T) {
+	for i := 0; i < len(slice)/2; i++ {
+		j := (len(slice) - 1) - i
+		slice[i], slice[j] = slice[j], slice[i]
+	}
+}
+
 func Max[T constraints.Ordered](a, b T) T {
 	if a > b {
 		return a
-- 
cgit v1.2.3-2-g168b