part: Optimize findIndex for node48

Use a 1-based indexing in node48.index so that 0 always means that
the key is not present.

Check for exact match first for node48 in findIndex, then check
if key is within bounds and only then binary search. Implement
binary search directly instead of using sort.Search to avoid
the function closure.

goos: darwin
goarch: arm64
pkg: github.com/cilium/statedb/part
cpu: Apple M1 Pro
                                   │     a.txt     │                b.txt                │
                                   │    sec/op     │    sec/op     vs base               │

_Uint64Map_Sequential_Insert-8        430.5µ ± ∞ ¹   476.4µ ± ∞ ¹        ~ (p=0.095 n=5)
_Uint64Map_Sequential_Txn_Insert-8    69.15µ ± ∞ ¹   68.59µ ± ∞ ¹        ~ (p=0.151 n=5)
_Uint64Map_Random_Insert-8            519.9µ ± ∞ ¹   541.2µ ± ∞ ¹        ~ (p=0.421 n=5)
_Uint64Map_Random_Txn_Insert-8        107.1µ ± ∞ ¹   104.0µ ± ∞ ¹   -2.87% (p=0.008 n=5)
_Insert_RootOnlyWatch-8               72.25µ ± ∞ ¹   69.49µ ± ∞ ¹        ~ (p=0.056 n=5)
_Insert-8                            100.18µ ± ∞ ¹   95.70µ ± ∞ ¹   -4.47% (p=0.016 n=5)
_GetInsert-8                          79.07µ ± ∞ ¹   78.45µ ± ∞ ¹        ~ (p=0.151 n=5)
_Hashmap_Insert-8                     36.09µ ± ∞ ¹   36.22µ ± ∞ ¹        ~ (p=0.841 n=5)
_Delete_Random-8                     10.188m ± ∞ ¹   8.640m ± ∞ ¹  -15.19% (p=0.008 n=5)
_find48-8                             2.038n ± ∞ ¹   2.065n ± ∞ ¹        ~ (p=0.095 n=5)
_findIndex48_hit-8                   10.180n ± ∞ ¹   2.047n ± ∞ ¹  -79.89% (p=0.008 n=5)
_findIndex48_miss-8                   8.468n ± ∞ ¹   2.260n ± ∞ ¹  -73.31% (p=0.008 n=5)

                        │    a.txt     │                b.txt                 │
                        │ objects/sec  │  objects/sec   vs base               │
_Insert_RootOnlyWatch-8   13.84M ± ∞ ¹    14.39M ± ∞ ¹        ~ (p=0.056 n=5)
_Insert-8                 9.982M ± ∞ ¹   10.449M ± ∞ ¹   +4.68% (p=0.016 n=5)
_GetInsert-8              12.65M ± ∞ ¹    12.75M ± ∞ ¹        ~ (p=0.151 n=5)
_Hashmap_Insert-8         27.71M ± ∞ ¹    27.61M ± ∞ ¹        ~ (p=0.841 n=5)
_Delete_Random-8          9.816M ± ∞ ¹   11.573M ± ∞ ¹  +17.91% (p=0.008 n=5)

Signed-off-by: Jussi Maki <jussi.maki@isovalent.com>

Author: joamaki

part/node.go

@@ -6,7 +6,6 @@ package part
 import (
 	"bytes"
 	"fmt"
-	"sort"
 	"strings"
 	"unsafe"
 )
@@ -246,7 +245,7 @@ func (n *header[T]) promote(txnID uint64) *header[T] {
 		node48.leaf = n.getLeaf()
 		copy(node48.children[:], node16.children[:node16.size()])
 		for i, k := range node16.keys[:node16.size()] {
-			node48.index[k] = int8(i)
+			node48.index[k] = uint8(i + 1)
 		}
 		if n.watch != nil {
 			node48.watch = make(chan struct{})
@@ -374,16 +373,34 @@ func (n *header[T]) findIndex(key byte) (*header[T], int) {
 		}
 		return nil, size
 	case nodeKind48:
-		children := n.children()
-		idx := sort.Search(len(children), func(i int) bool {
-			return children[i].prefix()[0] >= key
-		})
-		if idx >= n.size() || children[idx].prefix()[0] != key {
-			// No node found, return nil and the index into
-			// which it should go.
-			return nil, idx
+		n48 := n.node48()
+		// Check for exact match first
+		if idx := n48.index[key]; idx != 0 {
+			i := int(idx - 1)
+			return n48.children[i], i
+		}
+		// No exact match. Binary search to find insertion point.
+		size := n48.size()
+		children := n48.children[:size]
+		// Is the key between smallest and highest?
+		switch {
+		case key < children[0].key():
+			return nil, 0
+		case key > children[size-1].key():
+			return nil, size
 		}
-		return children[idx], idx
+		// No exact match, but key is in range. Binary search to find insertion point.
+		// We're not using sort.Search as that requires a function closure.
+		lo, hi := 0, size
+		for lo < hi {
+			mid := int(uint(lo+hi) / 2)
+			if children[mid].key() < key {
+				lo = mid + 1
+			} else {
+				hi = mid
+			}
+		}
+		return nil, lo
 	case nodeKind256:
 		return n.node256().children[key], int(key)
 	default:
@@ -418,10 +435,10 @@ func (n *header[T]) find(key byte) *header[T] {
 	case nodeKind48:
 		n48 := n.node48()
 		idx := n48.index[key]
-		if idx < 0 {
+		if idx == 0 {
 			return nil
 		}
-		return n48.children[idx]
+		return n48.children[idx-1]
 	case nodeKind256:
 		return n.node256().children[key]
 	default:
@@ -452,11 +469,11 @@ func (n *header[T]) insert(idx int, child *header[T]) {
 		n48 := n.node48()
 		for i := size - 1; i >= idx; i-- {
 			c := n48.children[i]
-			n48.index[c.key()] = int8(i + 1)
+			n48.index[c.key()] = uint8(i + 2)
 			n48.children[i+1] = c
 		}
 		n48.children[idx] = child
-		n48.index[child.key()] = int8(idx)
+		n48.index[child.key()] = uint8(idx + 1)
 	case nodeKind256:
 		n.node256().children[child.key()] = child
 	default:
@@ -489,9 +506,9 @@ func (n *header[T]) remove(idx int) {
 		for i := idx; i < newSize; i++ {
 			child := children[i+1]
 			children[i] = child
-			n48.index[child.key()] = int8(i)
+			n48.index[child.key()] = uint8(i + 1)
 		}
-		n48.index[key] = -1
+		n48.index[key] = 0
 		children[newSize] = nil
 	case nodeKind256:
 		n.node256().children[idx] = nil
@@ -547,8 +564,8 @@ type node48[T any] struct {
 	header[T]
 	txnID    uint64 // transaction ID that last mutated this node
 	children [48]*header[T]
-	leaf     *leaf[T] // non-nil if this node contains a value
-	index    [256]int8
+	leaf     *leaf[T]   // non-nil if this node contains a value
+	index    [256]uint8 // 1-based index for key in [children] (0 is absent, 1 is children[0])
 }
 
 type node256[T any] struct {

part/part_test.go

@@ -1667,20 +1667,17 @@ func Benchmark_find48(b *testing.B) {
 		header:   header[bool]{},
 		children: [48]*header[bool]{},
 		leaf:     nil,
-		index:    [256]int8{},
+		index:    [256]uint8{},
 	}
 	keyBytes := [48]byte{}
 	leaves := [48]leaf[bool]{}
-	for i := range n.index {
-		n.index[i] = -1
-	}
 	for i := range 48 {
 		keyBytes[i] = byte(i)
 		leaves[i].prefixP = &keyBytes[i]
 		leaves[i].prefixLen = 1
 		leaves[i].setKind(nodeKindLeaf)
 		n.children[i] = leaves[i].self()
-		n.index[byte(i)] = int8(i)
+		n.index[byte(i)] = uint8(i + 1)
 	}
 	n.setKind(nodeKind48)
 	n.setSize(48)
@@ -1695,7 +1692,7 @@ func Benchmark_findIndex48_hit(b *testing.B) {
 		header:   header[bool]{},
 		children: [48]*header[bool]{},
 		leaf:     nil,
-		index:    [256]int8{},
+		index:    [256]uint8{},
 	}
 	keyBytes := [48]byte{}
 	leaves := [48]leaf[bool]{}
@@ -1705,7 +1702,7 @@ func Benchmark_findIndex48_hit(b *testing.B) {
 		leaves[i].prefixLen = 1
 		leaves[i].setKind(nodeKindLeaf)
 		n.children[i] = leaves[i].self()
-		n.index[byte(i)] = int8(i + 1)
+		n.index[byte(i)] = uint8(i + 1)
 	}
 	n.setKind(nodeKind48)
 	n.setSize(48)
@@ -1720,20 +1717,17 @@ func Benchmark_findIndex48_miss(b *testing.B) {
 		header:   header[bool]{},
 		children: [48]*header[bool]{},
 		leaf:     nil,
-		index:    [256]int8{},
+		index:    [256]uint8{},
 	}
 	keyBytes := [48]byte{}
 	leaves := [48]leaf[bool]{}
-	for i := range n.index {
-		n.index[i] = -1
-	}
 	for i := range 48 {
 		keyBytes[i] = byte(i)
 		leaves[i].prefixP = &keyBytes[i]
 		leaves[i].prefixLen = 1
 		leaves[i].setKind(nodeKindLeaf)
 		n.children[i] = leaves[i].self()
-		n.index[byte(i)] = int8(i)
+		n.index[byte(i)] = uint8(i + 1)
 	}
 	n.setKind(nodeKind48)
 	n.setSize(48)

part/txn.go

@@ -596,15 +596,15 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[
 		demoted.setKind(nodeKind48)
 		demoted.setSize(size - 1)
 		demoted.setTxnID(txn.txnID)
-		n48 := demoted.node48()
-		n48.leaf = parent.getLeaf()
-		children := n48.children[:0]
-		for k, n := range parent.node256().children[:] {
-			if k != index && n != nil {
-				n48.index[k] = int8(len(children))
-				children = append(children, n)
+			n48 := demoted.node48()
+			n48.leaf = parent.getLeaf()
+			children := n48.children[:0]
+			for k, n := range parent.node256().children[:] {
+				if k != index && n != nil {
+					n48.index[k] = uint8(len(children) + 1)
+					children = append(children, n)
+				}
 			}
-		}
 		newParent = demoted
 	case parent.kind() == nodeKind48 && size <= 17:
 		demoted := (&node16[T]{header: *parent}).self()