Fixed a bug that when sequences were deleted and we cleaned up empty nodes we would not redo heights and balances.

This caused a rotate operation to possibly return nil and replace our root with nil when non empty.

Signed-off-by: Derek Collison <derek@nats.io>

server/avl/seqset.go

@@ -140,7 +140,7 @@ func (ss *SequenceSet) Heights() (l, r int) {
 
 // Returns min, max and number of set items.
 func (ss *SequenceSet) State() (min, max, num uint64) {
-	if ss.root == nil {
+	if ss == nil || ss.root == nil {
 		return 0, 0, 0
 	}
 	min, max = ss.MinMax()
@@ -446,20 +446,16 @@ func (n *node) insert(seq uint64, inserted *bool, nodes *int) *node {
 	// Don't make a function, impacts performance.
 	if bf := balanceF(n); bf > 1 {
 		// Left unbalanced.
-		if n.l.base+numEntries > seq {
-			return n.rotateR()
-		} else {
+		if balanceF(n.l) < 0 {
 			n.l = n.l.rotateL()
-			return n.rotateR()
 		}
+		return n.rotateR()
 	} else if bf < -1 {
-		// right unbalanced.
-		if n.r.base+numEntries > seq {
+		// Right unbalanced.
+		if balanceF(n.r) > 0 {
 			n.r = n.r.rotateR()
-			return n.rotateL()
-		} else {
-			return n.rotateL()
 		}
+		return n.rotateL()
 	}
 	return n
 }
@@ -507,6 +503,9 @@ func balanceF(n *node) int {
 }
 
 func maxH(n *node) int {
+	if n == nil {
+		return 0
+	}
 	var lh, rh int
 	if n.l != nil {
 		lh = n.l.h
@@ -550,39 +549,66 @@ func (n *node) delete(seq uint64, deleted *bool, nodes *int) *node {
 		n.r = n.r.delete(seq, deleted, nodes)
 	} else if empty := n.clear(seq, deleted); empty {
 		*nodes--
-		if nn := n.l; nn == nil {
+		if n.l == nil {
 			n = n.r
-		} else if nn.r == nil {
-			nn.r = n.r
-			n = nn
+		} else if n.r == nil {
+			n = n.l
 		} else {
-			nn.r.r = n.r
-			n = nn
+			// We have both children.
+			n.r = n.r.insertNodePrev(n.l)
+			n = n.r
 		}
 	}
 
+	if n != nil {
+		n.h = maxH(n) + 1
+	}
+
 	// Check balance.
 	if bf := balanceF(n); bf > 1 {
 		// Left unbalanced.
-		if n.l.base+numEntries > seq {
-			return n.rotateR()
-		} else {
+		if balanceF(n.l) < 0 {
 			n.l = n.l.rotateL()
-			return n.rotateR()
 		}
+		return n.rotateR()
 	} else if bf < -1 {
 		// right unbalanced.
-		if n.r.base+numEntries > seq {
+		if balanceF(n.r) > 0 {
 			n.r = n.r.rotateR()
-			return n.rotateL()
-		} else {
-			return n.rotateL()
 		}
+		return n.rotateL()
 	}
 
 	return n
 }
 
+// Will insert nn into the node assuming it is less than all other nodes in n.
+// Will re-calculate height and balance.
+func (n *node) insertNodePrev(nn *node) *node {
+	if n.l == nil {
+		n.l = nn
+	} else {
+		n.l = n.l.insertNodePrev(nn)
+	}
+	n.h = maxH(n) + 1
+
+	// Check balance.
+	if bf := balanceF(n); bf > 1 {
+		// Left unbalanced.
+		if balanceF(n.l) < 0 {
+			n.l = n.l.rotateL()
+		}
+		return n.rotateR()
+	} else if bf < -1 {
+		// right unbalanced.
+		if balanceF(n.r) > 0 {
+			n.r = n.r.rotateR()
+		}
+		return n.rotateL()
+	}
+	return n
+}
+
 func (n *node) exists(seq uint64) bool {
 	seq -= n.base
 	i := seq / bitsPerBucket

server/avl/seqset_test.go

@@ -134,21 +134,46 @@ func TestSeqSetDelete(t *testing.T) {
 		require_True(t, !ss.Exists(seq))
 	}
 	require_True(t, ss.root == nil)
+}
 
-	num := 22*numEntries + 22
+func TestSeqSetInsertAndDeletePedantic(t *testing.T) {
+	var ss SequenceSet
+
+	num := 50*numEntries + 22
 	nums := make([]uint64, 0, num)
 	for i := 0; i < num; i++ {
 		nums = append(nums, uint64(i))
 	}
 	rand.Shuffle(len(nums), func(i, j int) { nums[i], nums[j] = nums[j], nums[i] })
 
-	for _, n := range nums {
-		ss.Insert(n)
+	// Make sure always balanced.
+	testBalanced := func() {
+		t.Helper()
+		// Check heights.
+		ss.root.nodeIter(func(n *node) {
+			if n != nil && n.h != maxH(n)+1 {
+				t.Fatalf("Node height is wrong: %+v", n)
+			}
+		})
+		// Check balance factor.
+		if bf := balanceF(ss.root); bf > 1 || bf < -1 {
+			t.Fatalf("Unbalanced tree")
+		}
 	}
 
+	for _, n := range nums {
+		ss.Insert(n)
+		testBalanced()
+	}
+	require_True(t, ss.root != nil)
+
 	for _, n := range nums {
 		ss.Delete(n)
+		testBalanced()
 		require_True(t, !ss.Exists(n))
+		if ss.Size() > 0 {
+			require_True(t, ss.root != nil)
+		}
 	}
 	require_True(t, ss.root == nil)
 }