removed ItemSearchFn, simplified cursor usage

go/store/prolly/tree/map.go

@@ -80,12 +80,12 @@ func DiffKeyRangeOrderedTrees[K, V ~[]byte, O Ordering[K]](
 			return err
 		}
 	} else {
-		fromStart, err = NewCursorAtItem(ctx, from.NodeStore, from.Root, Item(start), from.searchNode)
+		fromStart, err = NewCursorAtKey(ctx, from.NodeStore, from.Root, start, from.Order)
 		if err != nil {
 			return err
 		}
 
-		toStart, err = NewCursorAtItem(ctx, to.NodeStore, to.Root, Item(start), to.searchNode)
+		toStart, err = NewCursorAtKey(ctx, to.NodeStore, to.Root, start, to.Order)
 		if err != nil {
 			return err
 		}
@@ -102,12 +102,12 @@ func DiffKeyRangeOrderedTrees[K, V ~[]byte, O Ordering[K]](
 			return err
 		}
 	} else {
-		fromStop, err = NewCursorAtItem(ctx, from.NodeStore, from.Root, Item(stop), from.searchNode)
+		fromStop, err = NewCursorAtKey(ctx, from.NodeStore, from.Root, stop, from.Order)
 		if err != nil {
 			return err
 		}
 
-		toStop, err = NewCursorAtItem(ctx, to.NodeStore, to.Root, Item(stop), to.searchNode)
+		toStop, err = NewCursorAtKey(ctx, to.NodeStore, to.Root, stop, to.Order)
 		if err != nil {
 			return err
 		}
@@ -180,7 +180,7 @@ func (t StaticMap[K, V, O]) WalkNodes(ctx context.Context, cb NodeCb) error {
 }
 
 func (t StaticMap[K, V, O]) Get(ctx context.Context, query K, cb KeyValueFn[K, V]) (err error) {
-	cur, err := NewLeafCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode)
+	cur, err := NewLeafCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order)
 	if err != nil {
 		return err
 	}
@@ -200,7 +200,7 @@ func (t StaticMap[K, V, O]) Get(ctx context.Context, query K, cb KeyValueFn[K, V
 }
 
 func (t StaticMap[K, V, O]) Has(ctx context.Context, query K) (ok bool, err error) {
-	cur, err := NewLeafCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode)
+	cur, err := NewLeafCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order)
 	if err != nil {
 		return false, err
 	}
@@ -389,7 +389,7 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi
 			return nil, nil, err
 		}
 	} else {
-		lo, err = NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(startInclusive), t.searchNode)
+		lo, err = NewCursorAtKey(ctx, t.NodeStore, t.Root, startInclusive, t.Order)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -401,7 +401,7 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi
 			return nil, nil, err
 		}
 	} else {
-		hi, err = NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(stopExclusive), t.searchNode)
+		hi, err = NewCursorAtKey(ctx, t.NodeStore, t.Root, stopExclusive, t.Order)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -410,35 +410,9 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi
 	return
 }
 
-// searchNode returns the smallest index where nd[i] >= query
-// Adapted from search.Sort to inline comparison.
-func (t StaticMap[K, V, O]) searchNode(query Item, nd Node) int {
-	n := int(nd.Count())
-	// Define f(-1) == false and f(n) == true.
-	// Invariant: f(i-1) == false, f(j) == true.
-	i, j := 0, n
-	for i < j {
-		h := int(uint(i+j) >> 1) // avoid overflow when computing h
-		less := t.Order.Compare(K(query), K(nd.GetKey(h))) <= 0
-		// i ≤ h < j
-		if !less {
-			i = h + 1 // preserves f(i-1) == false
-		} else {
-			j = h // preserves f(j) == true
-		}
-	}
-	// i == j, f(i-1) == false, and
-	// f(j) (= f(i)) == true  =>  answer is i.
-	return i
-}
-
-func (t StaticMap[K, V, O]) CompareItems(left, right Item) int {
-	return t.Order.Compare(K(left), K(right))
-}
-
 // getOrdinalForKey returns the smallest ordinal position at which the key >= |query|.
 func (t StaticMap[K, V, O]) GetOrdinalForKey(ctx context.Context, query K) (uint64, error) {
-	cur, err := NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode)
+	cur, err := NewCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order)
 	if err != nil {
 		return 0, err
 	}
@@ -446,8 +420,6 @@ func (t StaticMap[K, V, O]) GetOrdinalForKey(ctx context.Context, query K) (uint
 	return GetOrdinalOfCursor(cur)
 }
 
-var _ ItemSearchFn = StaticMap[Item, Item, Ordering[Item]]{}.searchNode
-
 type OrderedTreeIter[K, V ~[]byte] struct {
 	// current tuple location
 	curr *Cursor

go/store/prolly/tree/node_cursor.go

@@ -43,31 +43,6 @@ type Ordering[K ~[]byte] interface {
 	Compare(left, right K) int
 }
 
-// SearchForKey returns a SearchFn for |key|.
-func SearchForKey[K ~[]byte, O Ordering[K]](key K, order O) SearchFn {
-	return func(nd Node) (idx int) {
-		n := int(nd.Count())
-		// Define f(-1) == false and f(n) == true.
-		// Invariant: f(i-1) == false, f(j) == true.
-		i, j := 0, n
-		for i < j {
-			h := int(uint(i+j) >> 1) // avoid overflow when computing h
-			less := order.Compare(key, K(nd.GetKey(h))) <= 0
-			// i ≤ h < j
-			if !less {
-				i = h + 1 // preserves f(i-1) == false
-			} else {
-				j = h // preserves f(j) == true
-			}
-		}
-		// i == j, f(i-1) == false, and
-		// f(j) (= f(i)) == true  =>  answer is i.
-		return i
-	}
-}
-
-type ItemSearchFn func(item Item, nd Node) (idx int)
-
 func NewCursorAtStart(ctx context.Context, ns NodeStore, nd Node) (cur *Cursor, err error) {
 	cur = &Cursor{nd: nd, nrw: ns}
 	var leaf bool
@@ -240,49 +215,29 @@ func NewCursorFromSearchFn(ctx context.Context, ns NodeStore, nd Node, search Se
 	return
 }
 
-func NewCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item, search ItemSearchFn) (cur *Cursor, err error) {
-	cur = &Cursor{nd: nd, nrw: ns}
+func NewLeafCursorAtKey[K ~[]byte, O Ordering[K]](ctx context.Context, ns NodeStore, nd Node, key K, order O) (Cursor, error) {
+	cur := Cursor{nd: nd, nrw: ns}
+	for {
+		// binary search |cur.nd| for |key|
+		i, j := 0, cur.nd.Count()
+		for i < j {
+			h := int(uint(i+j) >> 1)
+			cmp := order.Compare(key, K(cur.nd.GetKey(h)))
+			if cmp > 0 {
+				i = h + 1
+			} else {
+				j = h
+			}
+		}
+		cur.idx = i
 
-	cur.idx = search(item, cur.nd)
-	var leaf bool
-	leaf, err = cur.isLeaf()
-	if err != nil {
-		return nil, err
-	}
-	for !leaf {
-
-		// stay in bounds for internal nodes
-		cur.keepInBounds()
-
-		nd, err = fetchChild(ctx, ns, cur.CurrentRef())
+		leaf, err := cur.isLeaf()
 		if err != nil {
 			return cur, err
+		} else if leaf {
+			break // done
 		}
 
-		parent := cur
-		cur = &Cursor{nd: nd, parent: parent, nrw: ns}
-
-		cur.idx = search(item, cur.nd)
-		leaf, err = cur.isLeaf()
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	return
-}
-
-func NewLeafCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item, search ItemSearchFn) (cur Cursor, err error) {
-	cur = Cursor{nd: nd, parent: nil, nrw: ns}
-
-	cur.idx = search(item, cur.nd)
-	var leaf bool
-	leaf, err = cur.isLeaf()
-	if err != nil {
-		return cur, err
-	}
-	for !leaf {
-
 		// stay in bounds for internal nodes
 		cur.keepInBounds()
 
@@ -291,17 +246,33 @@ func NewLeafCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item,
 		if err != nil {
 			return cur, err
 		}
-
-		cur.idx = search(item, cur.nd)
-		leaf, err = cur.isLeaf()
-		if err != nil {
-			return cur, err
-		}
 	}
-
 	return cur, nil
 }
 
+// SearchForKey returns a SearchFn for |key|.
+func SearchForKey[K ~[]byte, O Ordering[K]](key K, order O) SearchFn {
+	return func(nd Node) (idx int) {
+		n := int(nd.Count())
+		// Define f(-1) == false and f(n) == true.
+		// Invariant: f(i-1) == false, f(j) == true.
+		i, j := 0, n
+		for i < j {
+			h := int(uint(i+j) >> 1) // avoid overflow when computing h
+			less := order.Compare(key, K(nd.GetKey(h))) <= 0
+			// i ≤ h < j
+			if !less {
+				i = h + 1 // preserves f(i-1) == false
+			} else {
+				j = h // preserves f(j) == true
+			}
+		}
+		// i == j, f(i-1) == false, and
+		// f(j) (= f(i)) == true  =>  answer is i.
+		return i
+	}
+}
+
 type LeafSpan struct {
 	Leaves     []Node
 	LocalStart int

go/store/prolly/tree/node_cursor_test.go

@@ -17,7 +17,6 @@ package tree
 import (
 	"context"
 	"fmt"
-	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -79,7 +78,7 @@ func testNewCursorAtItem(t *testing.T, count int) {
 	ctx := context.Background()
 	for i := range items {
 		key, value := items[i][0], items[i][1]
-		cur, err := NewCursorAtItem(ctx, ns, root, key, searchTestTree)
+		cur, err := NewCursorAtKey(ctx, ns, root, val.Tuple(key), keyDesc)
 		require.NoError(t, err)
 		assert.Equal(t, key, cur.CurrentKey())
 		assert.Equal(t, value, cur.CurrentValue())
@@ -89,18 +88,11 @@ func testNewCursorAtItem(t *testing.T, count int) {
 }
 
 func testGetOrdinalOfCursor(t *testing.T, count int) {
-	tuples, d := AscendingUintTuples(count)
-
-	search := func(item Item, nd Node) (idx int) {
-		return sort.Search(int(nd.count), func(i int) bool {
-			l, r := val.Tuple(item), val.Tuple(nd.GetKey(i))
-			return d.Compare(l, r) <= 0
-		})
-	}
+	tuples, desc := AscendingUintTuples(count)
 
 	ctx := context.Background()
 	ns := NewTestNodeStore()
-	serializer := message.NewProllyMapSerializer(d, ns.Pool())
+	serializer := message.NewProllyMapSerializer(desc, ns.Pool())
 	chkr, err := newEmptyChunker(ctx, ns, serializer)
 	require.NoError(t, err)
 
@@ -112,7 +104,7 @@ func testGetOrdinalOfCursor(t *testing.T, count int) {
 	assert.NoError(t, err)
 
 	for i := 0; i < len(tuples); i++ {
-		curr, err := NewCursorAtItem(ctx, ns, nd, Item(tuples[i][0]), search)
+		curr, err := NewCursorAtKey(ctx, ns, nd, tuples[i][0], desc)
 		require.NoError(t, err)
 
 		ord, err := GetOrdinalOfCursor(curr)
@@ -121,11 +113,11 @@ func testGetOrdinalOfCursor(t *testing.T, count int) {
 		assert.Equal(t, uint64(i), ord)
 	}
 
-	b := val.NewTupleBuilder(d)
+	b := val.NewTupleBuilder(desc)
 	b.PutUint32(0, uint32(len(tuples)))
 	aboveItem := b.Build(sharedPool)
 
-	curr, err := NewCursorAtItem(ctx, ns, nd, Item(aboveItem), search)
+	curr, err := NewCursorAtKey(ctx, ns, nd, aboveItem, desc)
 	require.NoError(t, err)
 
 	ord, err := GetOrdinalOfCursor(curr)
@@ -171,13 +163,6 @@ var valDesc = val.NewTupleDescriptor(
 	val.Type{Enc: val.Int64Enc, Nullable: true},
 )
 
-func searchTestTree(item Item, nd Node) int {
-	return sort.Search(int(nd.count), func(i int) bool {
-		l, r := val.Tuple(item), val.Tuple(nd.GetKey(i))
-		return keyDesc.Compare(l, r) <= 0
-	})
-}
-
 func randomTupleItemPairs(count int, ns NodeStore) (items [][2]Item) {
 	tups := RandomTuplePairs(count, keyDesc, valDesc, ns)
 	items = make([][2]Item, count)