optimize iterate over loose N48 node

Signed-off-by: you06 <you1474600@gmail.com>

Author: you06

tree_benchmark_test.go

@@ -189,3 +189,27 @@ func BenchmarkHSKTreeForEach(b *testing.B) {
 	stats := collectStats(tree.Iterator(TraverseAll))
 	assert.Equal(b, treeStats{4995, 1630, 276, 21, 4}, stats)
 }
+
+func BenchmarkN48LooseScan(b *testing.B) {
+	tree := New()
+
+	for i := 0; i < 17; i++ {
+		for j := 0; j < 17; j++ {
+			for k := 0; k < 17; k++ {
+				for l := 0; l < 17; l++ {
+					key := []byte{byte(i), byte(j), byte(k), byte(l)}
+					tree.Insert(key, key)
+				}
+			}
+		}
+	}
+
+	subBench := func(b *testing.B, options ...int) {
+		b.ResetTimer()
+		stats := collectStats(tree.Iterator(options...))
+		assert.Equal(b, treeStats{83521, 0, 0, 5220, 0}, stats)
+	}
+
+	b.Run("Iterator", func(b *testing.B) { subBench(b, TraverseAll) })
+	b.Run("IteratorReverse", func(b *testing.B) { subBench(b, TraverseAll, TraverseReverse) })
+}

tree_traversal.go

@@ -1,5 +1,10 @@
 package art
 
+import (
+	"math"
+	"math/bits"
+)
+
 // traverseAction is an action to be taken during tree traversal.
 type traverseAction int
 
@@ -104,28 +109,40 @@ func (ctx *traverse48Context) ascTraversal() (int, bool) {
 		return node48Max, true
 	}
 
-	for ; ctx.curKeyIdx < node256Max; ctx.curKeyIdx++ {
-		if ctx.n48.hasChild(ctx.curKeyIdx) {
+	for curOffset := ctx.curKeyIdx >> n48bitShift; curOffset < 4; curOffset++ {
+		offset := ctx.curKeyIdx % n48maskLen
+		mask := math.MaxUint64 - (uint64(1) << offset) + 1 // e.g. offset=3 => 0b111...111000
+		curr := ctx.n48.present[curOffset] & mask
+		zeros := bits.TrailingZeros64(curr)
+		if zeros < n48maskLen {
+			ctx.curKeyIdx = curOffset*n48maskLen + zeros
 			ctx.curKeyCh = ctx.n48.keys[ctx.curKeyIdx]
 			ctx.curKeyIdx++
-
 			return int(ctx.curKeyCh), true
 		}
+		ctx.curKeyIdx = 0
 	}
+	ctx.curKeyIdx = node256Max
 
 	return 0, false
 }
 
 // descTraversal traverses the children in descending order.
 func (ctx *traverse48Context) descTraversal() (int, bool) {
-	for ; ctx.curKeyIdx > 0; ctx.curKeyIdx-- {
-		if ctx.n48.hasChild(ctx.curKeyIdx) {
+	for curOffset := ctx.curKeyIdx >> n48bitShift; curOffset >= 0; curOffset-- {
+		offset := ctx.curKeyIdx % n48maskLen
+		mask := uint64(1)<<(offset+1) - 1 // e.g. offset=3 => 0b000...0001111
+		curr := ctx.n48.present[curOffset] & mask
+		zeros := bits.LeadingZeros64(curr)
+		if zeros < n48maskLen {
+			ctx.curKeyIdx = curOffset*n48maskLen + n48maskLen - (zeros + 1)
 			ctx.curKeyCh = ctx.n48.keys[ctx.curKeyIdx]
 			ctx.curKeyIdx--
-
 			return int(ctx.curKeyCh), true
 		}
+		ctx.curKeyIdx = n48maskLen - 1
 	}
+	ctx.curKeyIdx = -1
 
 	if !ctx.zeroChildDone {
 		ctx.zeroChildDone = true

tree_traversal_test.go

@@ -589,3 +589,40 @@ func TestIteratorEmptyTreeBehavior(t *testing.T) {
 	assert.Nil(t, n)
 	assert.Equal(t, ErrNoMoreNodes, err)
 }
+
+func TestIteratorStatsWithReverse(t *testing.T) {
+	t.Parallel()
+
+	testCases := map[int]Kind{
+		node4Max + 1:  Node16,
+		node16Max + 1: Node48,
+		node48Max + 1: Node256,
+	}
+	for leafCount, expectedKind := range testCases {
+		tree := New()
+		for i := 0; i < leafCount; i++ {
+			key := []byte{byte(i)}
+			tree.Insert(key, key)
+		}
+
+		expectedStats := treeStats{
+			leafCount: leafCount,
+		}
+		switch expectedKind {
+		case Node16:
+			expectedStats.node16Count = 1
+		case Node48:
+			expectedStats.node48Count = 1
+		case Node256:
+			expectedStats.node256Count = 1
+		default:
+			t.Fatalf("unexpected expectedKind: %v", expectedKind)
+		}
+
+		stats := collectStats(tree.Iterator(TraverseAll))
+		assert.Equal(t, expectedStats, stats)
+
+		stats = collectStats(tree.Iterator(TraverseAll, TraverseReverse))
+		assert.Equal(t, expectedStats, stats)
+	}
+}