diff options
| -rw-r--r-- | pkg/rbtree/range.go | 23 | ||||
| -rw-r--r-- | pkg/util/generic.go | 7 |
2 files changed, 30 insertions, 0 deletions
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 |