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/rbtree/rbtree.go | 522 --------------------- lib/rbtree/rbtree_test.go | 171 ------- ...841300ea7e6bac1a1e9505b1535810083d18db95d86f489 | 2 - ...b25d402bffbfdce114baad0b1c8b9a9189b5a97fda27e41 | 2 - 4 files changed, 697 deletions(-) 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/rbtree') 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