diff options
Diffstat (limited to 'pkg/rbtree/rbtree_test.go')
-rw-r--r-- | pkg/rbtree/rbtree_test.go | 134 |
1 files changed, 134 insertions, 0 deletions
diff --git a/pkg/rbtree/rbtree_test.go b/pkg/rbtree/rbtree_test.go new file mode 100644 index 0000000..6c49127 --- /dev/null +++ b/pkg/rbtree/rbtree_test.go @@ -0,0 +1,134 @@ +package rbtree + +import ( + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/constraints" +) + +func printTree[K constraints.Ordered, V any](t *testing.T, tree *Tree[K, V]) { + fmtBareNode := func(node *Node[V]) string { + if node == nil { + return "nil" + } + if node.Color == Red { + return fmt.Sprintf("R(%v)", node.Value) + } else { + return fmt.Sprintf("B(%v)", node.Value) + } + } + addIndent := func(indent string, lines []string) []string { + ret := make([]string, 0, len(lines)) + for _, line := range lines { + ret = append(ret, indent+line) + } + return ret + } + var fmtNode func(node *Node[V]) []string + fmtNode = func(node *Node[V]) []string { + if node == nil { + return []string{"nil"} + } + ret := addIndent(" ", fmtNode(node.Right)) + ret = append(ret, fmtBareNode(node)) + ret = append(ret, addIndent(" ", fmtNode(node.Left))...) + return ret + } + t.Log("\n" + strings.Join(fmtNode(tree.root), "\n")) +} + +func checkTree[K constraints.Ordered, V any](t *testing.T, tree *Tree[K, V]) { + // 1. Every node is either red or black + + // 2. The root is black. + assert.Equal(t, Black, tree.root.getColor()) + + // 3. Every nil is black. + + // 4. If a node is red, then both its children are black. + tree.Walk(func(node *Node[V]) { + if node.getColor() == Red { + assert.Equal(t, Black, node.Left.getColor()) + assert.Equal(t, Black, node.Right.getColor()) + } + }) + + // 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) + } + tree.Walk(func(node *Node[V]) { + var cnts []int + walkCnt(node, 0, func(cnt int) { + cnts = append(cnts, cnt) + }) + for i := range cnts { + if cnts[0] != cnts[i] { + if !assert.Truef(t, false, "node %v: not all leafs have same black-count: %v", node.Value, cnts) { + printTree(t, tree) + } + } + } + }) +} + +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[uint8, uint8]{ + KeyFn: func(x uint8) uint8 { return x }, + } + checkTree(t, tree) + for _, b := range dat { + ins := (b & 0b0100_0000) != 0 + val := (b & 0b0011_1111) + if ins { + t.Logf("Insert(%v)", val) + tree.Insert(val) + node := tree.Lookup(val) + require.NotNil(t, node) + assert.Equal(t, val, node.Value) + } else { + t.Logf("Delete(%v)", val) + tree.Delete(val) + assert.Nil(t, tree.Lookup(val)) + } + checkTree(t, tree) + } + }) +} |