net/art: move child table pointers out of strideEntry

In preparation for a different refactor, but incidentally also saves
10-25% memory on overall table size in benchmarks.

Updates #7781

Signed-off-by: David Anderson <danderson@tailscale.com>

3 files changed, 50 insertions, 49 deletions

diff --git a/net/art/stride_table.go b/net/art/stride_table.go
index 79f607c77..ea261efac 100644
--- a/net/art/stride_table.go
+++ b/net/art/stride_table.go
@@ -29,8 +29,6 @@ type strideEntry[T any] struct {
 	prefixIndex int
 	// value is the value associated with the strideEntry, if any.
 	value *T
-	// child is the child strideTable associated with the strideEntry, if any.
-	child *strideTable[T]
 }
 
 // strideTable is a binary tree that implements an 8-bit routing table.
@@ -50,12 +48,17 @@ type strideTable[T any] struct {
 	// parent of the node at index i is located at index i>>1, and its children
 	// at indices i<<1 and (i<<1)+1.
 	//
-	// A few consequences of this arrangement: host routes (/8) occupy the last
-	// 256 entries in the table; the single default route /0 is at index 1, and
-	// index 0 is unused (in the original paper, it's hijacked through sneaky C
-	// memory trickery to store the refcount, but this is Go, where we don't
-	// store random bits in pointers lest we confuse the GC)
+	// A few consequences of this arrangement: host routes (/8) occupy
+	// the last numChildren entries in the table; the single default
+	// route /0 is at index 1, and index 0 is unused (in the original
+	// paper, it's hijacked through sneaky C memory trickery to store
+	// the refcount, but this is Go, where we don't store random bits
+	// in pointers lest we confuse the GC)
 	entries [lastHostIndex + 1]strideEntry[T]
+	// children are the child tables of this table. Each child
+	// represents the address space within one of this table's host
+	// routes (/8).
+	children [numChildren]*strideTable[T]
 	// routeRefs is the number of route entries in this table.
 	routeRefs uint16
 	// childRefs is the number of child strideTables referenced by this table.
@@ -67,63 +70,60 @@ const (
 	firstHostIndex = 0b1_0000_0000
 	// lastHostIndex is the array index of the last host route. This is hostIndex(0xFF/8).
 	lastHostIndex = 0b1_1111_1111
+
+	// numChildren is the maximum number of child tables a strideTable can hold.
+	numChildren = 256
 )
 
-// getChild returns the child strideTable pointer for addr (if any), and an
-// internal array index that can be used with deleteChild.
-func (t *strideTable[T]) getChild(addr uint8) (child *strideTable[T], idx int) {
-	idx = hostIndex(addr)
-	return t.entries[idx].child, idx
+// getChild returns the child strideTable pointer for addr, or nil if none.
+func (t *strideTable[T]) getChild(addr uint8) *strideTable[T] {
+	return t.children[addr]
 }
 
-// deleteChild deletes the child strideTable at idx (if any). idx should be
-// obtained via a call to getChild.
-func (t *strideTable[T]) deleteChild(idx int) {
-	t.entries[idx].child = nil
-	t.childRefs--
+// deleteChild deletes the child strideTable at addr. It is valid to
+// delete a non-existent child.
+func (t *strideTable[T]) deleteChild(addr uint8) {
+	if t.children[addr] != nil {
+		t.childRefs--
+	}
+	t.children[addr] = nil
 }
 
-// setChild replaces the child strideTable for addr (if any) with child.
+// setChild sets the child strideTable for addr to child.
 func (t *strideTable[T]) setChild(addr uint8, child *strideTable[T]) {
-	t.setChildByIndex(hostIndex(addr), child)
-}
-
-// setChildByIndex replaces the child strideTable at idx (if any) with
-// child. idx should be obtained via a call to getChild.
-func (t *strideTable[T]) setChildByIndex(idx int, child *strideTable[T]) {
-	if t.entries[idx].child == nil {
+	if t.children[addr] == nil {
 		t.childRefs++
 	}
-	t.entries[idx].child = child
+	t.children[addr] = child
 }
 
 // getOrCreateChild returns the child strideTable for addr, creating it if
 // necessary.
 func (t *strideTable[T]) getOrCreateChild(addr uint8) (child *strideTable[T], created bool) {
-	idx := hostIndex(addr)
-	if t.entries[idx].child == nil {
-		t.entries[idx].child = &strideTable[T]{
+	ret := t.children[addr]
+	if ret == nil {
+		ret = &strideTable[T]{
 			prefix: childPrefixOf(t.prefix, addr),
 		}
+		t.children[addr] = ret
 		t.childRefs++
-		return t.entries[idx].child, true
+		return ret, true
 	}
-	return t.entries[idx].child, false
+	return ret, false
 }
 
 // getValAndChild returns both the prefix and child strideTable for
 // addr. Both returned values can be nil if no entry of that type
 // exists for addr.
 func (t *strideTable[T]) getValAndChild(addr uint8) (*T, *strideTable[T]) {
-	idx := hostIndex(addr)
-	return t.entries[idx].value, t.entries[idx].child
+	return t.entries[hostIndex(addr)].value, t.children[addr]
 }
 
 // findFirstChild returns the first child strideTable in t, or nil if
 // t has no children.
 func (t *strideTable[T]) findFirstChild() *strideTable[T] {
-	for i := firstHostIndex; i <= lastHostIndex; i++ {
-		if child := t.entries[i].child; child != nil {
+	for _, child := range t.children {
+		if child != nil {
 			return child
 		}
 	}
diff --git a/net/art/table.go b/net/art/table.go
index 42123704d..4b12d4bd7 100644
--- a/net/art/table.go
+++ b/net/art/table.go
@@ -364,7 +364,7 @@ func (t *Table[T]) Delete(pfx netip.Prefix) {
 	// write to strideTables[N] and strideIndexes[N-1].
 	strideIdx := 0
 	strideTables := [16]*strideTable[T]{st}
-	strideIndexes := [15]int{}
+	strideIndexes := [15]uint8{}
 
 	// Similar to Insert, navigate down the tree of strideTables,
 	// looking for the one that houses this prefix. This part is
@@ -384,7 +384,7 @@ func (t *Table[T]) Delete(pfx netip.Prefix) {
 		if debugDelete {
 			fmt.Printf("delete: loop byteIdx=%d numBits=%d st.prefix=%s\n", byteIdx, numBits, st.prefix)
 		}
-		child, idx := st.getChild(bs[byteIdx])
+		child := st.getChild(bs[byteIdx])
 		if child == nil {
 			// Prefix can't exist in the table, because one of the
 			// necessary strideTables doesn't exist.
@@ -393,7 +393,7 @@ func (t *Table[T]) Delete(pfx netip.Prefix) {
 			}
 			return
 		}
-		strideIndexes[strideIdx] = idx
+		strideIndexes[strideIdx] = bs[byteIdx]
 		strideTables[strideIdx+1] = child
 		strideIdx++
 
@@ -475,7 +475,7 @@ func (t *Table[T]) Delete(pfx netip.Prefix) {
 			if debugDelete {
 				fmt.Printf("delete: compact parent.prefix=%s st.prefix=%s child.prefix=%s\n", parent.prefix, cur.prefix, child.prefix)
 			}
-			strideTables[strideIdx-1].setChildByIndex(strideIndexes[strideIdx-1], child)
+			strideTables[strideIdx-1].setChild(strideIndexes[strideIdx-1], child)
 			return
 		default:
 			// This table has two or more children, so it's acting as a "fork in
@@ -505,12 +505,12 @@ func strideSummary[T any](w io.Writer, st *strideTable[T], indent int) {
 	fmt.Fprintf(w, "%s: %d routes, %d children\n", st.prefix, st.routeRefs, st.childRefs)
 	indent += 4
 	st.treeDebugStringRec(w, 1, indent)
-	for i := firstHostIndex; i <= lastHostIndex; i++ {
-		if child := st.entries[i].child; child != nil {
-			addr, len := inversePrefixIndex(i)
-			fmt.Fprintf(w, "%s%d/%d (%02x/%d): ", strings.Repeat(" ", indent), addr, len, addr, len)
-			strideSummary(w, child, indent)
+	for addr, child := range st.children {
+		if child == nil {
+			continue
 		}
+		fmt.Fprintf(w, "%s%d/8 (%02x/8): ", strings.Repeat(" ", indent), addr, addr)
+		strideSummary(w, child, indent)
 	}
 }
 
diff --git a/net/art/table_test.go b/net/art/table_test.go
index 0885b3c02..835828340 100644
--- a/net/art/table_test.go
+++ b/net/art/table_test.go
@@ -607,7 +607,7 @@ func TestInsertCompare(t *testing.T) {
 			seenVals4[fastVal] = true
 		}
 		if slowVal != fastVal {
-			t.Errorf("get(%q) = %p, want %p", a, fastVal, slowVal)
+			t.Fatalf("get(%q) = %p, want %p", a, fastVal, slowVal)
 		}
 	}
 
@@ -1092,11 +1092,12 @@ func (t *Table[T]) numStridesRec(seen map[*strideTable[T]]bool, st *strideTable[
 	if st.childRefs == 0 {
 		return ret
 	}
-	for i := firstHostIndex; i <= lastHostIndex; i++ {
-		if c := st.entries[i].child; c != nil && !seen[c] {
-			seen[c] = true
-			ret += t.numStridesRec(seen, c)
+	for _, c := range st.children {
+		if c == nil || seen[c] {
+			continue
 		}
+		seen[c] = true
+		ret += t.numStridesRec(seen, c)
 	}
 	return ret
 }