perf: galloping search in sealed LID iterators' NextGeq

frac/sealed/lids/gallop.go

@@ -0,0 +1,66 @@
+package lids
+
+// Galloping (exponential) searches for the NextGeq hot path: zigzag targets
+// are monotone and the chunk slice narrows from the consumption edge, so the
+// answer is usually a few elements away from the edge. Probing exponentially
+// from the edge costs ~2*log2(d) for an answer d away (bounded by ~2x a full
+// binary search over the remainder, which pays log2(len) regardless of d).
+
+// searchGeqGallop returns the index of the first element >= t, probing from
+// the front (IteratorDesc consumes the slice front-to-back).
+func searchGeqGallop(a []uint32, t uint32) int {
+	n := len(a)
+	if n == 0 || a[0] >= t {
+		return 0
+	}
+	i := 1
+	for i < n && a[i] < t {
+		i <<= 1
+	}
+	// a[i>>1] < t, so the answer is in (i>>1, min(i, n)].
+	lo := i>>1 + 1
+	hi := i
+	if hi > n {
+		hi = n
+	}
+	for lo < hi {
+		m := int(uint(lo+hi) >> 1)
+		if a[m] < t {
+			lo = m + 1
+		} else {
+			hi = m
+		}
+	}
+	return lo
+}
+
+// searchGtGallopTail returns the index of the first element > t, probing from
+// the back (IteratorAsc consumes the slice back-to-front).
+func searchGtGallopTail(a []uint32, t uint32) int {
+	n := len(a)
+	if n == 0 {
+		return 0
+	}
+	if a[n-1] <= t {
+		return n
+	}
+	i := 1
+	for i < n && a[n-1-i] > t {
+		i <<= 1
+	}
+	// a[n-1-(i>>1)] > t, so the answer is <= n-1-(i>>1) < hi.
+	hi := n - i>>1
+	lo := n - i
+	if lo < 0 {
+		lo = 0
+	}
+	for lo < hi {
+		m := int(uint(lo+hi) >> 1)
+		if a[m] <= t {
+			lo = m + 1
+		} else {
+			hi = m
+		}
+	}
+	return lo
+}

frac/sealed/lids/gallop_test.go

@@ -0,0 +1,31 @@
+package lids
+
+import (
+	"math/rand/v2"
+	"sort"
+	"testing"
+)
+
+func TestGallopAgainstSortSearch(t *testing.T) {
+	r := rand.New(rand.NewPCG(3, 4))
+	for i := 0; i < 200000; i++ {
+		n := 1 + r.IntN(200)
+		a := make([]uint32, n)
+		v := uint32(r.IntN(3))
+		for j := range a {
+			v += uint32(r.IntN(4)) // duplicates allowed
+			a[j] = v
+		}
+		q := uint32(r.IntN(int(v) + 3))
+
+		wantGeq := sort.Search(len(a), func(k int) bool { return a[k] >= q })
+		if got := searchGeqGallop(a, q); got != wantGeq {
+			t.Fatalf("searchGeqGallop(%v, %d) = %d, want %d", a, q, got, wantGeq)
+		}
+
+		wantGt := sort.Search(len(a), func(k int) bool { return a[k] > q })
+		if got := searchGtGallopTail(a, q); got != wantGt {
+			t.Fatalf("searchGtGallopTail(%v, %d) = %d, want %d", a, q, got, wantGt)
+		}
+	}
+}

frac/sealed/lids/iterator_asc.go

@@ -93,7 +93,7 @@ func (it *IteratorAsc) NextGeq(nextID node.LID) node.LID {
 			continue
 		}
 
-		idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] > nextID.Unpack() }) - 1
+		idx := searchGtGallopTail(it.lids, nextID.Unpack()) - 1
 		if idx >= 0 {
 			lid := it.lids[idx]
 			it.lids = it.lids[:idx]
@@ -126,7 +126,7 @@ func (it *IteratorAsc) NextBatchGeq(nextID node.LID) node.LIDBatch {
 			continue
 		}
 
-		idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] > nextID.Unpack() }) - 1
+		idx := searchGtGallopTail(it.lids, nextID.Unpack()) - 1
 		if idx >= 0 {
 			batch := it.lids[:idx+1]
 			it.lids = it.lids[:0]

frac/sealed/lids/iterator_desc.go

@@ -92,7 +92,7 @@ func (it *IteratorDesc) NextGeq(nextID node.LID) node.LID {
 			continue
 		}
 
-		idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] >= nextID.Unpack() })
+		idx := searchGeqGallop(it.lids, nextID.Unpack())
 		if idx < len(it.lids) {
 			it.lids = it.lids[idx:]
 			lid := it.lids[0]
@@ -130,7 +130,7 @@ func (it *IteratorDesc) NextBatchGeq(nextID node.LID) node.LIDBatch {
 			continue
 		}
 
-		idx := sort.Search(len(it.lids), func(i int) bool { return it.lids[i] >= nextID.Unpack() })
+		idx := searchGeqGallop(it.lids, nextID.Unpack())
 		if idx < len(it.lids) {
 			batch := it.lids[idx:len(it.lids)]
 			it.lids = it.lids[:0]