From 4e29bb393ec774f0a79c70d9d69c54fe4e8ecb72 Mon Sep 17 00:00:00 2001 From: Luke Shumaker Date: Wed, 13 Jul 2022 20:36:27 -0600 Subject: Move lib/rbtree to lib/containers --- lib/btrfs/btrfsvol/lvm.go | 22 +- lib/btrfsprogs/btrfsutil/broken_btree.go | 10 +- lib/containers/rbtree.go | 522 +++++++++++++++++++++ lib/containers/rbtree_test.go | 171 +++++++ ...841300ea7e6bac1a1e9505b1535810083d18db95d86f489 | 2 + ...b25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 | 2 + lib/rbtree/rbtree.go | 522 --------------------- lib/rbtree/rbtree_test.go | 171 ------- ...841300ea7e6bac1a1e9505b1535810083d18db95d86f489 | 2 - ...b25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 | 2 - 10 files changed, 713 insertions(+), 713 deletions(-) create mode 100644 lib/containers/rbtree.go create mode 100644 lib/containers/rbtree_test.go create mode 100644 lib/containers/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 create mode 100644 lib/containers/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 delete mode 100644 lib/rbtree/rbtree.go delete mode 100644 lib/rbtree/rbtree_test.go delete mode 100644 lib/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 delete mode 100644 lib/rbtree/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 (limited to 'lib') diff --git a/lib/btrfs/btrfsvol/lvm.go b/lib/btrfs/btrfsvol/lvm.go index ca46ad3..351cbe5 100644 --- a/lib/btrfs/btrfsvol/lvm.go +++ b/lib/btrfs/btrfsvol/lvm.go @@ -10,7 +10,7 @@ import ( "os" "reflect" - "git.lukeshu.com/btrfs-progs-ng/lib/rbtree" + "git.lukeshu.com/btrfs-progs-ng/lib/containers" "git.lukeshu.com/btrfs-progs-ng/lib/util" ) @@ -19,8 +19,8 @@ type LogicalVolume[PhysicalVolume util.File[PhysicalAddr]] struct { id2pv map[DeviceID]PhysicalVolume - logical2physical *rbtree.Tree[util.NativeOrdered[LogicalAddr], chunkMapping] - physical2logical map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping] + logical2physical *containers.RBTree[util.NativeOrdered[LogicalAddr], chunkMapping] + physical2logical map[DeviceID]*containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping] } var _ util.File[LogicalAddr] = (*LogicalVolume[util.File[PhysicalAddr]])(nil) @@ -30,18 +30,18 @@ func (lv *LogicalVolume[PhysicalVolume]) init() { lv.id2pv = make(map[DeviceID]PhysicalVolume) } if lv.logical2physical == nil { - lv.logical2physical = &rbtree.Tree[util.NativeOrdered[LogicalAddr], chunkMapping]{ + lv.logical2physical = &containers.RBTree[util.NativeOrdered[LogicalAddr], chunkMapping]{ KeyFn: func(chunk chunkMapping) util.NativeOrdered[LogicalAddr] { return util.NativeOrdered[LogicalAddr]{Val: chunk.LAddr} }, } } if lv.physical2logical == nil { - lv.physical2logical = make(map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping], len(lv.id2pv)) + lv.physical2logical = make(map[DeviceID]*containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping], len(lv.id2pv)) } for devid := range lv.id2pv { if _, ok := lv.physical2logical[devid]; !ok { - lv.physical2logical[devid] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{ + lv.physical2logical[devid] = &containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping]{ KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] { return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr} }, @@ -74,7 +74,7 @@ func (lv *LogicalVolume[PhysicalVolume]) AddPhysicalVolume(id DeviceID, dev Phys lv, dev.Name(), other.Name(), id) } lv.id2pv[id] = dev - lv.physical2logical[id] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{ + lv.physical2logical[id] = &containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping]{ KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] { return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr} }, @@ -173,8 +173,8 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error { } func (lv *LogicalVolume[PhysicalVolume]) fsck() error { - physical2logical := make(map[DeviceID]*rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]) - if err := lv.logical2physical.Walk(func(node *rbtree.Node[chunkMapping]) error { + physical2logical := make(map[DeviceID]*containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping]) + if err := lv.logical2physical.Walk(func(node *containers.RBNode[chunkMapping]) error { chunk := node.Value for _, stripe := range chunk.PAddrs { if _, devOK := lv.id2pv[stripe.Dev]; !devOK { @@ -182,7 +182,7 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error { lv, stripe.Dev) } if _, exists := physical2logical[stripe.Dev]; !exists { - physical2logical[stripe.Dev] = &rbtree.Tree[util.NativeOrdered[PhysicalAddr], devextMapping]{ + physical2logical[stripe.Dev] = &containers.RBTree[util.NativeOrdered[PhysicalAddr], devextMapping]{ KeyFn: func(ext devextMapping) util.NativeOrdered[PhysicalAddr] { return util.NativeOrdered[PhysicalAddr]{Val: ext.PAddr} }, @@ -216,7 +216,7 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error { func (lv *LogicalVolume[PhysicalVolume]) Mappings() []Mapping { var ret []Mapping - _ = lv.logical2physical.Walk(func(node *rbtree.Node[chunkMapping]) error { + _ = lv.logical2physical.Walk(func(node *containers.RBNode[chunkMapping]) error { chunk := node.Value var flags *BlockGroupFlags if chunk.Flags != nil { diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go index 2835b35..d2dc81c 100644 --- a/lib/btrfsprogs/btrfsutil/broken_btree.go +++ b/lib/btrfsprogs/btrfsutil/broken_btree.go @@ -14,7 +14,7 @@ import ( "git.lukeshu.com/btrfs-progs-ng/lib/btrfs" "git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol" - "git.lukeshu.com/btrfs-progs-ng/lib/rbtree" + "git.lukeshu.com/btrfs-progs-ng/lib/containers" "git.lukeshu.com/btrfs-progs-ng/lib/util" ) @@ -24,7 +24,7 @@ type indexItem struct { } type cachedIndex struct { - Index *rbtree.Tree[btrfs.Key, indexItem] + Index *containers.RBTree[btrfs.Key, indexItem] Err error } @@ -67,7 +67,7 @@ func NewBrokenTrees(ctx context.Context, inner *btrfs.FS) btrfs.Trees { } } -func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) (*rbtree.Tree[btrfs.Key, indexItem], error) { +func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) (*containers.RBTree[btrfs.Key, indexItem], error) { var treeRoot *btrfs.TreeRoot var err error if treeID == btrfs.ROOT_TREE_OBJECTID { @@ -92,7 +92,7 @@ func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) (*rbtree.Tree[btrfs.Key, in if err != nil { cacheEntry.Err = err } else { - cacheEntry.Index = &rbtree.Tree[btrfs.Key, indexItem]{ + cacheEntry.Index = &containers.RBTree[btrfs.Key, indexItem]{ KeyFn: func(item indexItem) btrfs.Key { return item.Key }, @@ -196,7 +196,7 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfs.ObjID, errHand return } var node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node] - _ = index.Walk(func(indexItem *rbtree.Node[indexItem]) error { + _ = index.Walk(func(indexItem *containers.RBNode[indexItem]) error { if ctx.Err() != nil { return ctx.Err() } 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 +// +// 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 +// +// 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") diff --git a/lib/rbtree/rbtree.go b/lib/rbtree/rbtree.go deleted file mode 100644 index 25f3b1c..0000000 --- a/lib/rbtree/rbtree.go +++ /dev/null @@ -1,522 +0,0 @@ -// Copyright (C) 2022 Luke Shumaker -// -// SPDX-License-Identifier: GPL-2.0-or-later - -package rbtree - -import ( - "fmt" - "reflect" - - "git.lukeshu.com/btrfs-progs-ng/lib/util" -) - -type Color bool - -const ( - Black = Color(false) - Red = Color(true) -) - -type Node[V any] struct { - Parent, Left, Right *Node[V] - - Color Color - - Value V -} - -func (node *Node[V]) getColor() Color { - if node == nil { - return Black - } - return node.Color -} - -type Tree[K util.Ordered[K], V any] struct { - KeyFn func(V) K - root *Node[V] -} - -func (t *Tree[K, V]) Walk(fn func(*Node[V]) error) error { - return t.root.walk(fn) -} - -func (node *Node[V]) walk(fn func(*Node[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 *Tree[K, V]) Search(fn func(V) int) *Node[V] { - ret, _ := t.root.search(fn) - return ret -} - -func (node *Node[V]) search(fn func(V) int) (exact, nearest *Node[V]) { - var prev *Node[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 *Tree[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 *Tree[K, V]) Lookup(key K) *Node[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 *Tree[K, V]) Min() *Node[V] { - return t.root.min() -} - -func (node *Node[V]) min() *Node[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 *Tree[K, V]) Max() *Node[V] { - return t.root.max() -} - -func (node *Node[V]) max() *Node[V] { - if node == nil { - return nil - } - for { - if node.Right == nil { - return node - } - node = node.Right - } -} - -func (t *Tree[K, V]) Next(cur *Node[V]) *Node[V] { - return cur.next() -} - -func (cur *Node[V]) next() *Node[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 *Tree[K, V]) Prev(cur *Node[V]) *Node[V] { - return cur.prev() -} - -func (cur *Node[V]) prev() *Node[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 *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: - 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 *Tree[K, V]) leftRotate(x *Node[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 *Tree[K, V]) rightRotate(y *Node[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 *Tree[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 := &Node[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 *Tree[K, V]) transplant(old, new *Node[V]) { - *t.parentChild(old) = new - if new != nil { - new.Parent = old.Parent - } -} - -func (t *Tree[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 *Node[V] - var nodeToRebalanceParent *Node[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/rbtree/rbtree_test.go b/lib/rbtree/rbtree_test.go deleted file mode 100644 index e6007e1..0000000 --- a/lib/rbtree/rbtree_test.go +++ /dev/null @@ -1,171 +0,0 @@ -// Copyright (C) 2022 Luke Shumaker -// -// SPDX-License-Identifier: GPL-2.0-or-later - -package rbtree - -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 *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 - - // 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 *Node[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 *Node[V], cnt int, leafFn func(int)) - walkCnt = func(node *Node[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 *Node[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 *Node[V]) error { - actOrder = append(actOrder, tree.KeyFn(node.Value).Val) - return nil - })) - require.Equal(t, expectedOrder, actOrder) -} - -func FuzzTree(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 := &Tree[util.NativeOrdered[uint8], uint8]{ - KeyFn: func(x uint8) util.NativeOrdered[uint8] { - return util.NativeOrdered[uint8]{Val: x} - }, - } - set := make(map[uint8]struct{}) - checkTree(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})) - } - checkTree(t, set, tree) - } - }) -} diff --git a/lib/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 b/lib/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 deleted file mode 100644 index 40318c6..0000000 --- a/lib/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 +++ /dev/null @@ -1,2 +0,0 @@ -go test fuzz v1 -[]byte("aAm0BCrb0x00!0000000000000") diff --git a/lib/rbtree/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 b/lib/rbtree/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 deleted file mode 100644 index 238e44f..0000000 --- a/lib/rbtree/testdata/fuzz/FuzzTree/f9e6421dacf921f7bb25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 +++ /dev/null @@ -1,2 +0,0 @@ -go test fuzz v1 -[]byte("YZAB\x990") -- cgit v1.2.3-2-g168b