diff options
Diffstat (limited to 'lib/containers')
4 files changed, 697 insertions, 0 deletions
diff --git a/lib/containers/rbtree.go b/lib/containers/rbtree.go new file mode 100644 index 0000000..6bffc02 --- /dev/null +++ b/lib/containers/rbtree.go @@ -0,0 +1,522 @@ +// Copyright (C) 2022 Luke Shumaker <lukeshu@lukeshu.com> +// +// SPDX-License-Identifier: GPL-2.0-or-later + +package containers + +import ( + "fmt" + "reflect" + + "git.lukeshu.com/btrfs-progs-ng/lib/util" +) + +type Color bool + +const ( + Black = Color(false) + Red = Color(true) +) + +type RBNode[V any] struct { + Parent, Left, Right *RBNode[V] + + Color Color + + Value V +} + +func (node *RBNode[V]) getColor() Color { + if node == nil { + return Black + } + return node.Color +} + +type RBTree[K util.Ordered[K], V any] struct { + KeyFn func(V) K + root *RBNode[V] +} + +func (t *RBTree[K, V]) Walk(fn func(*RBNode[V]) error) error { + return t.root.walk(fn) +} + +func (node *RBNode[V]) walk(fn func(*RBNode[V]) error) error { + if node == nil { + 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 + } + return nil +} + +// Search the tree for a value that satisfied the given callbackk +// function. A return value of 0 means to to return this value; <0 +// means to go left on the tree (the value is too high), >0 means to +// go right on th etree (the value is too low). +// +// +-----+ +// | v=8 | == 0 : this is it +// +-----+ +// / \ +// / \ +// <0 : go left >0 : go right +// / \ +// +---+ +---+ +// | 7 | | 9 | +// +---+ +---+ +// +// Returns nil if no such value is found. +// +// Search is good for advanced lookup, like when a range of values is +// acceptable. For simple exact-value lookup, use Lookup. +func (t *RBTree[K, V]) Search(fn func(V) int) *RBNode[V] { + ret, _ := t.root.search(fn) + return ret +} + +func (node *RBNode[V]) search(fn func(V) int) (exact, nearest *RBNode[V]) { + var prev *RBNode[V] + for { + if node == nil { + return nil, prev + } + direction := fn(node.Value) + prev = node + switch { + case direction < 0: + node = node.Left + case direction == 0: + return node, nil + case direction > 0: + node = node.Right + } + } +} + +func (t *RBTree[K, V]) exactKey(key K) func(V) int { + return func(val V) int { + valKey := t.KeyFn(val) + return key.Cmp(valKey) + } +} + +// Lookup looks up the value for an exact key. If no such value +// exists, nil is returned. +func (t *RBTree[K, V]) Lookup(key K) *RBNode[V] { + return t.Search(t.exactKey(key)) +} + +// Min returns the minimum value stored in the tree, or nil if the +// tree is empty. +func (t *RBTree[K, V]) Min() *RBNode[V] { + return t.root.min() +} + +func (node *RBNode[V]) min() *RBNode[V] { + if node == nil { + return nil + } + for { + if node.Left == nil { + return node + } + node = node.Left + } +} + +// Max returns the maximum value stored in the tree, or nil if the +// tree is empty. +func (t *RBTree[K, V]) Max() *RBNode[V] { + return t.root.max() +} + +func (node *RBNode[V]) max() *RBNode[V] { + if node == nil { + return nil + } + for { + if node.Right == nil { + return node + } + node = node.Right + } +} + +func (t *RBTree[K, V]) Next(cur *RBNode[V]) *RBNode[V] { + return cur.next() +} + +func (cur *RBNode[V]) next() *RBNode[V] { + if cur.Right != nil { + return cur.Right.min() + } + child, parent := cur, cur.Parent + for parent != nil && child == parent.Right { + child, parent = parent, parent.Parent + } + return parent +} + +func (t *RBTree[K, V]) Prev(cur *RBNode[V]) *RBNode[V] { + return cur.prev() +} + +func (cur *RBNode[V]) prev() *RBNode[V] { + if cur.Left != nil { + return cur.Left.max() + } + child, parent := cur, cur.Parent + for parent != nil && child == parent.Left { + child, parent = parent, parent.Parent + } + return parent +} + +// SearchRange is like Search, but returns all nodes that match the +// function; assuming that they are contiguous. +func (t *RBTree[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 *RBTree[K, V]) Equal(u *RBTree[K, V]) bool { + if (t == nil) != (u == nil) { + return false + } + if t == nil { + return true + } + + var tSlice []V + _ = t.Walk(func(node *RBNode[V]) error { + tSlice = append(tSlice, node.Value) + return nil + }) + + var uSlice []V + _ = u.Walk(func(node *RBNode[V]) error { + uSlice = append(uSlice, node.Value) + return nil + }) + + return reflect.DeepEqual(tSlice, uSlice) +} + +func (t *RBTree[K, V]) parentChild(node *RBNode[V]) **RBNode[V] { + switch { + case node.Parent == nil: + return &t.root + case node.Parent.Left == node: + return &node.Parent.Left + case node.Parent.Right == node: + return &node.Parent.Right + default: + panic(fmt.Errorf("node %p is not a child of its parent %p", node, node.Parent)) + } +} + +func (t *RBTree[K, V]) leftRotate(x *RBNode[V]) { + // p p + // | | + // +---+ +---+ + // | x | | y | + // +---+ +---+ + // / \ => / \ + // a +---+ +---+ c + // | y | | x | + // +---+ +---+ + // / \ / \ + // b c a b + + // Define 'p', 'x', 'y', and 'b' per the above diagram. + p := x.Parent + pChild := t.parentChild(x) + y := x.Right + b := y.Left + + // Move things around + + y.Parent = p + *pChild = y + + x.Parent = y + y.Left = x + + if b != nil { + b.Parent = x + } + x.Right = b +} + +func (t *RBTree[K, V]) rightRotate(y *RBNode[V]) { + // | | + // +---+ +---+ + // | y | | x | + // +---+ +---+ + // / \ => / \ + // +---+ c a +---+ + // | x | | y | + // +---+ +---+ + // / \ / \ + // a b b c + + // Define 'p', 'x', 'y', and 'b' per the above diagram. + p := y.Parent + pChild := t.parentChild(y) + x := y.Left + b := x.Right + + // Move things around + + x.Parent = p + *pChild = x + + y.Parent = x + x.Right = y + + if b != nil { + b.Parent = y + } + y.Left = b +} + +func (t *RBTree[K, V]) Insert(val V) { + // Naive-insert + + key := t.KeyFn(val) + exact, parent := t.root.search(t.exactKey(key)) + if exact != nil { + exact.Value = val + return + } + + node := &RBNode[V]{ + Color: Red, + Parent: parent, + Value: val, + } + if parent == nil { + t.root = node + } else if key.Cmp(t.KeyFn(parent.Value)) < 0 { + parent.Left = node + } else { + parent.Right = node + } + + // Re-balance + // + // This is closely based on the algorithm presented in CLRS + // 3e. + + for node.Parent.getColor() == Red { + if node.Parent == node.Parent.Parent.Left { + uncle := node.Parent.Parent.Right + if uncle.getColor() == Red { + node.Parent.Color = Black + uncle.Color = Black + node.Parent.Parent.Color = Red + node = node.Parent.Parent + } else { + if node == node.Parent.Right { + node = node.Parent + t.leftRotate(node) + } + node.Parent.Color = Black + node.Parent.Parent.Color = Red + t.rightRotate(node.Parent.Parent) + } + } else { + uncle := node.Parent.Parent.Left + if uncle.getColor() == Red { + node.Parent.Color = Black + uncle.Color = Black + node.Parent.Parent.Color = Red + node = node.Parent.Parent + } else { + if node == node.Parent.Left { + node = node.Parent + t.rightRotate(node) + } + node.Parent.Color = Black + node.Parent.Parent.Color = Red + t.leftRotate(node.Parent.Parent) + } + } + } + t.root.Color = Black +} + +func (t *RBTree[K, V]) transplant(old, new *RBNode[V]) { + *t.parentChild(old) = new + if new != nil { + new.Parent = old.Parent + } +} + +func (t *RBTree[K, V]) Delete(key K) { + nodeToDelete := t.Lookup(key) + if nodeToDelete == nil { + return + } + + // This is closely based on the algorithm presented in CLRS + // 3e. + + var nodeToRebalance *RBNode[V] + var nodeToRebalanceParent *RBNode[V] // in case 'nodeToRebalance' is nil, which it can be + needsRebalance := nodeToDelete.Color == Black + + switch { + case nodeToDelete.Left == nil: + nodeToRebalance = nodeToDelete.Right + nodeToRebalanceParent = nodeToDelete.Parent + t.transplant(nodeToDelete, nodeToDelete.Right) + case nodeToDelete.Right == nil: + nodeToRebalance = nodeToDelete.Left + nodeToRebalanceParent = nodeToDelete.Parent + t.transplant(nodeToDelete, nodeToDelete.Left) + default: + // The node being deleted has a child on both sides, + // so we've go to reshuffle the parents a bit to make + // room for those children. + next := nodeToDelete.next() + if next.Parent == nodeToDelete { + // p p + // | | + // +-----+ +-----+ + // | ntd | | nxt | + // +-----+ +-----+ + // / \ => / \ + // a +-----+ a b + // | nxt | + // +-----+ + // / \ + // nil b + nodeToRebalance = next.Right + nodeToRebalanceParent = next + + *t.parentChild(nodeToDelete) = next + next.Parent = nodeToDelete.Parent + + next.Left = nodeToDelete.Left + next.Left.Parent = next + } else { + // p p + // | | + // +-----+ +-----+ + // | ntd | | nxt | + // +-----+ +-----+ + // / \ / \ + // a x a x + // / \ => / \ + // y z y z + // / \ / \ + // +-----+ c b c + // | nxt | + // +-----+ + // / \ + // nil b + y := next.Parent + b := next.Right + nodeToRebalance = b + nodeToRebalanceParent = y + + *t.parentChild(nodeToDelete) = next + next.Parent = nodeToDelete.Parent + + next.Left = nodeToDelete.Left + next.Left.Parent = next + + next.Right = nodeToDelete.Right + next.Right.Parent = next + + y.Left = b + if b != nil { + b.Parent = y + } + } + + // idk + needsRebalance = next.Color == Black + next.Color = nodeToDelete.Color + } + + if needsRebalance { + node := nodeToRebalance + nodeParent := nodeToRebalanceParent + for node != t.root && node.getColor() == Black { + if node == nodeParent.Left { + sibling := nodeParent.Right + if sibling.getColor() == Red { + sibling.Color = Black + nodeParent.Color = Red + t.leftRotate(nodeParent) + sibling = nodeParent.Right + } + if sibling.Left.getColor() == Black && sibling.Right.getColor() == Black { + sibling.Color = Red + node, nodeParent = nodeParent, nodeParent.Parent + } else { + if sibling.Right.getColor() == Black { + sibling.Left.Color = Black + sibling.Color = Red + t.rightRotate(sibling) + sibling = nodeParent.Right + } + sibling.Color = nodeParent.Color + nodeParent.Color = Black + sibling.Right.Color = Black + t.leftRotate(nodeParent) + node, nodeParent = t.root, nil + } + } else { + sibling := nodeParent.Left + if sibling.getColor() == Red { + sibling.Color = Black + nodeParent.Color = Red + t.rightRotate(nodeParent) + sibling = nodeParent.Left + } + if sibling.Right.getColor() == Black && sibling.Left.getColor() == Black { + sibling.Color = Red + node, nodeParent = nodeParent, nodeParent.Parent + } else { + if sibling.Left.getColor() == Black { + sibling.Right.Color = Black + sibling.Color = Red + t.leftRotate(sibling) + sibling = nodeParent.Left + } + sibling.Color = nodeParent.Color + nodeParent.Color = Black + sibling.Left.Color = Black + t.rightRotate(nodeParent) + node, nodeParent = t.root, nil + } + } + } + if node != nil { + node.Color = Black + } + } +} diff --git a/lib/containers/rbtree_test.go b/lib/containers/rbtree_test.go new file mode 100644 index 0000000..3360bc0 --- /dev/null +++ b/lib/containers/rbtree_test.go @@ -0,0 +1,171 @@ +// Copyright (C) 2022 Luke Shumaker <lukeshu@lukeshu.com> +// +// SPDX-License-Identifier: GPL-2.0-or-later + +package containers + +import ( + "fmt" + "io" + "sort" + "strings" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/exp/constraints" + + "git.lukeshu.com/btrfs-progs-ng/lib/util" +) + +func (t *RBTree[K, V]) ASCIIArt() string { + var out strings.Builder + t.root.asciiArt(&out, "", "", "") + return out.String() +} + +func (node *RBNode[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 *RBNode[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 checkRBTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]struct{}, tree *RBTree[util.NativeOrdered[K], V]) { + // 1. Every node is either red or black + + // 2. The root is black. + require.Equal(t, Black, tree.root.getColor()) + + // 3. Every nil is black. + + // 4. If a node is red, then both its children are black. + require.NoError(t, tree.Walk(func(node *RBNode[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 + // nodes. + var walkCnt func(node *RBNode[V], cnt int, leafFn func(int)) + walkCnt = func(node *RBNode[V], cnt int, leafFn func(int)) { + if node.getColor() == Black { + cnt++ + } + if node == nil { + leafFn(cnt) + return + } + walkCnt(node.Left, cnt, leafFn) + walkCnt(node.Right, cnt, leafFn) + } + require.NoError(t, tree.Walk(func(node *RBNode[V]) error { + var cnts []int + walkCnt(node, 0, func(cnt int) { + cnts = append(cnts, cnt) + }) + for i := range cnts { + if cnts[0] != cnts[i] { + require.Truef(t, false, "node %v: not all leafs have same black-count: %v", node.Value, cnts) + break + } + } + return nil + })) + + // expected contents + expectedOrder := make([]K, 0, len(expectedSet)) + for k := range expectedSet { + expectedOrder = append(expectedOrder, k) + node := tree.Lookup(util.NativeOrdered[K]{Val: k}) + require.NotNil(t, tree, node) + require.Equal(t, k, tree.KeyFn(node.Value).Val) + } + sort.Slice(expectedOrder, func(i, j int) bool { + return expectedOrder[i] < expectedOrder[j] + }) + actOrder := make([]K, 0, len(expectedSet)) + require.NoError(t, tree.Walk(func(node *RBNode[V]) error { + actOrder = append(actOrder, tree.KeyFn(node.Value).Val) + return nil + })) + require.Equal(t, expectedOrder, actOrder) +} + +func FuzzRBTree(f *testing.F) { + Ins := uint8(0b0100_0000) + Del := uint8(0) + + f.Add([]uint8{}) + f.Add([]uint8{Ins | 5, Del | 5}) + f.Add([]uint8{Ins | 5, Del | 6}) + f.Add([]uint8{Del | 6}) + + f.Add([]uint8{ // CLRS Figure 14.4 + Ins | 1, + Ins | 2, + Ins | 5, + Ins | 7, + Ins | 8, + Ins | 11, + Ins | 14, + Ins | 15, + + Ins | 4, + }) + + f.Fuzz(func(t *testing.T, dat []uint8) { + tree := &RBTree[util.NativeOrdered[uint8], uint8]{ + KeyFn: func(x uint8) util.NativeOrdered[uint8] { + return util.NativeOrdered[uint8]{Val: x} + }, + } + set := make(map[uint8]struct{}) + checkRBTree(t, set, tree) + t.Logf("\n%s\n", tree.ASCIIArt()) + for _, b := range dat { + ins := (b & 0b0100_0000) != 0 + val := (b & 0b0011_1111) + if ins { + t.Logf("Insert(%v)", val) + tree.Insert(val) + set[val] = struct{}{} + t.Logf("\n%s\n", tree.ASCIIArt()) + node := tree.Lookup(util.NativeOrdered[uint8]{Val: val}) + require.NotNil(t, node) + require.Equal(t, val, node.Value) + } else { + t.Logf("Delete(%v)", val) + tree.Delete(util.NativeOrdered[uint8]{Val: val}) + delete(set, val) + t.Logf("\n%s\n", tree.ASCIIArt()) + require.Nil(t, tree.Lookup(util.NativeOrdered[uint8]{Val: val})) + } + checkRBTree(t, set, tree) + } + }) +} diff --git a/lib/containers/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 b/lib/containers/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 new file mode 100644 index 0000000..40318c6 --- /dev/null +++ b/lib/containers/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 @@ -0,0 +1,2 @@ +go test fuzz v1 +[]byte("aAm0BCrb0x00!0000000000000") diff --git a/lib/containers/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 b/lib/containers/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 new file mode 100644 index 0000000..238e44f --- /dev/null +++ b/lib/containers/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 @@ -0,0 +1,2 @@ +go test fuzz v1 +[]byte("YZAB\x990") |