diff options
Diffstat (limited to 'lib/rbtree')
-rw-r--r-- | lib/rbtree/print_test.go | 45 | ||||
-rw-r--r-- | lib/rbtree/rbtree.go | 42 | ||||
-rw-r--r-- | lib/rbtree/rbtree_test.go | 37 | ||||
-rw-r--r-- | lib/rbtree/rbtree_util.go | 52 |
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) -} |