fix2

4 files changed, 69 insertions, 64 deletions

diff --git a/pkg/rbtree/print_test.go b/pkg/rbtree/print_test.go
index c4154c8..3e37cf2 100644
--- a/pkg/rbtree/print_test.go
+++ b/pkg/rbtree/print_test.go
@@ -12,6 +12,17 @@ func (t *Tree[K, V]) ASCIIArt() string {
 	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)
@@ -21,9 +32,9 @@ func (node *Node[V]) asciiArt(w io.Writer, u, m, l string) {
 	node.Right.asciiArt(w, u+"     ", u+"  ,--", u+"  |  ")
 
 	if node.Color == Red {
-		fmt.Fprintf(w, "%sR(%v)\n", m, node.Value)
+		fmt.Fprintf(w, "%s%v\n", m, node)
 	} else {
-		fmt.Fprintf(w, "%sB(%v)\n", m, node.Value)
+		fmt.Fprintf(w, "%s%v\n", m, node)
 	}
 
 	node.Left.asciiArt(w, l+"  |  ", l+"  `--", l+"     ")
diff --git a/pkg/rbtree/rbtree.go b/pkg/rbtree/rbtree.go
index 34c73d4..5cc9ac6 100644
--- a/pkg/rbtree/rbtree.go
+++ b/pkg/rbtree/rbtree.go
@@ -330,93 +330,86 @@ func (t *Tree[K, V]) Delete(key K) {
 		return
 	}
 
-	// A pointer to node that now resides at the place in the tree
-	// where the deleted node was.
-	nodeToRebalance := t.parentChild(nodeToDelete)
-	nodeToRebalanceParent := nodeToDelete.Parent
+	var nodeToRebalance, nodeToRebalanceParent *Node[V]
 	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 nodeToDelete.next().Parent == nodeToDelete, then
-		// this is pretty easy...
-		if next.Parent != nodeToDelete {
-			// but if it's not, then we have another step
-			// ('mid' might actually be a chain of nodes)
-			// to get it there
-			//
+		if next.Parent == nodeToDelete {
 			//         p                  p
 			//         |                  |
 			//      +-----+            +-----+
-			//      | ntd |            | ntd |
+			//      | ntd |            | nxt |
 			//      +-----+            +-----+
-			//      /     \            /     \
-			//     a       x          a    +-----+
-			//            / \     =>       | nxt |
-			//           y   z             +-----+
-			//          / \                /     \
-			//    +-----+  c              nil     x
-			//    | nxt |                        / \
-			//    +-----+                       y   z
-			//    /     \                      / \
-			//  nil     b                     b   c
-			//
-			// This looks an aweful lot like a
-			// t.rightRotate(next.Parent), but spans
-			// multiple nodes
-			x := nodeToDelete.Right
+			//      /     \       =>   /     \
+			//     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.Parent = nodeToDelete
-			nodeToDelete.Right = next
+			next.Left = nodeToDelete.Left
+			next.Left.Parent = next
 
-			x.Parent = next
-			next.Right = x
+			next.Right = nodeToDelete.Right
+			next.Right.Parent = next
 
+			y.Left = b
 			if b != nil {
 				b.Parent = y
 			}
-			y.Left = b
 		}
-		// ... OK, back to the easy case:
-		//
-		//         p                  p
-		//         |                  |
-		//      +-----+            +-----+
-		//      | ntd |            | nxt |
-		//      +-----+            +-----+
-		//      /     \       =>   /     \
-		//     a     +-----+      a      b
-		//           | nxt |
-		//           +-----+
-		//            /   \
-		//          nil   b
-		//
-		*t.parentChild(nodeToDelete) = next
-		next.Parent = nodeToDelete.Parent
-
-		next.Left = nodeToDelete.Left
-		next.Left.Parent = next
 
 		// idk
-		nodeToRebalance = &next.Right
-		nodeToRebalanceParent = next
 		needsRebalance = next.Color == Black
 		next.Color = nodeToDelete.Color
 	}
 
 	if needsRebalance {
-		node := *nodeToRebalance
+		node := nodeToRebalance
 		nodeParent := nodeToRebalanceParent // in case 'node' is nil, which it can be
 		for node != t.root && node.getColor() == Black {
 			if node == nodeParent.Left {
diff --git a/pkg/rbtree/rbtree_test.go b/pkg/rbtree/rbtree_test.go
index 4f4a12c..15db778 100644
--- a/pkg/rbtree/rbtree_test.go
+++ b/pkg/rbtree/rbtree_test.go
@@ -3,7 +3,6 @@ package rbtree
 import (
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/constraints"
 )
@@ -12,15 +11,15 @@ 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())
+	require.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())
+			require.Equal(t, Black, node.Left.getColor())
+			require.Equal(t, Black, node.Right.getColor())
 		}
 	})
 
@@ -46,7 +45,7 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, tree *Tree[K, V]) {
 		})
 		for i := range cnts {
 			if cnts[0] != cnts[i] {
-				assert.Truef(t, false, "node %v: not all leafs have same black-count: %v", node.Value, cnts)
+				require.Truef(t, false, "node %v: not all leafs have same black-count: %v", node.Value, cnts)
 				break
 			}
 		}
@@ -90,12 +89,12 @@ func FuzzTree(f *testing.F) {
 				t.Logf("\n%s\n", tree.ASCIIArt())
 				node := tree.Lookup(val)
 				require.NotNil(t, node)
-				assert.Equal(t, val, node.Value)
+				require.Equal(t, val, node.Value)
 			} else {
 				t.Logf("Delete(%v)", val)
 				tree.Delete(val)
 				t.Logf("\n%s\n", tree.ASCIIArt())
-				assert.Nil(t, tree.Lookup(val))
+				require.Nil(t, tree.Lookup(val))
 			}
 			checkTree(t, tree)
 		}
diff --git a/pkg/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489 b/pkg/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489
new file mode 100644
index 0000000..40318c6
--- /dev/null
+++ b/pkg/rbtree/testdata/fuzz/FuzzTree/be408ce7760bc8ced841300ea7e6bac1a1e9505b1535810083d18db95d86f489
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("aAm0BCrb0x00!0000000000000")