net/art: allow non-pointers as values

Values are still turned into pointers internally to maintain the
invariants of strideTable, but from the user's perspective it's
now possible to tbl.Insert(pfx, true) rather than
tbl.Insert(pfx, ptr.To(true)).

Updates tailscale#7781

Signed-off-by: David Anderson <[email protected]>

Author: danderson

net/art/stride_table.go

@@ -18,19 +18,6 @@ const (
 	debugStrideDelete = false
 )
 
-// strideEntry is a strideTable entry.
-type strideEntry[T any] struct {
-	// prefixIndex is the prefixIndex(...) value that caused this stride entry's
-	// value to be populated, or 0 if value is nil.
-	//
-	// We need to keep track of this because allot() uses it to determine
-	// whether an entry was propagated from a parent entry, or if it's a
-	// different independent route.
-	prefixIndex int
-	// value is the value associated with the strideEntry, if any.
-	value *T
-}
-
 // strideTable is a binary tree that implements an 8-bit routing table.
 //
 // The leaves of the binary tree are host routes (/8s). Each parent is a
@@ -54,7 +41,9 @@ type strideTable[T any] struct {
 	// 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]
+	//
+	// A nil value means no route matches the queried route.
+	entries [lastHostIndex + 1]*T
 	// children are the child tables of this table. Each child
 	// represents the address space within one of this table's host
 	// routes (/8).
@@ -112,13 +101,6 @@ func (t *strideTable[T]) getOrCreateChild(addr uint8) (child *strideTable[T], cr
 	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]) {
-	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] {
@@ -130,73 +112,115 @@ func (t *strideTable[T]) findFirstChild() *strideTable[T] {
 	return nil
 }
 
+// hasPrefixRootedAt reports whether t.entries[idx] is the root node of
+// a prefix.
+func (t *strideTable[T]) hasPrefixRootedAt(idx int) bool {
+	val := t.entries[idx]
+	if val == nil {
+		return false
+	}
+
+	parentIdx := parentIndex(idx)
+	if parentIdx == 0 {
+		// idx is non-nil, and is at the 0/0 route position.
+		return true
+	}
+	if parent := t.entries[parentIdx]; val != parent {
+		// parent node in the tree isn't the same prefix, so idx must
+		// be a root.
+		return true
+	}
+	return false
+}
+
 // allot updates entries whose stored prefixIndex matches oldPrefixIndex, in the
 // subtree rooted at idx. Matching entries have their stored prefixIndex set to
 // newPrefixIndex, and their value set to val.
 //
 // allot is the core of the ART algorithm, enabling efficient insertion/deletion
 // while preserving very fast lookups.
-func (t *strideTable[T]) allot(idx int, oldPrefixIndex, newPrefixIndex int, val *T) {
-	if t.entries[idx].prefixIndex != oldPrefixIndex {
-		// current prefixIndex isn't what we expect. This is a recursive call
-		// that found a child subtree that already has a more specific route
-		// installed. Don't touch it.
+func (t *strideTable[T]) allot(idx int, old, new *T) {
+	if t.entries[idx] != old {
+		// current idx isn't what we expect. This is a recursive call
+		// that found a child subtree that already has a more specific
+		// route installed. Don't touch it.
 		return
 	}
-	t.entries[idx].value = val
-	t.entries[idx].prefixIndex = newPrefixIndex
+	t.entries[idx] = new
 	if idx >= firstHostIndex {
 		// The entry we just updated was a host route, we're at the bottom of
 		// the binary tree.
 		return
 	}
 	// Propagate the allotment to this node's children.
 	left := idx << 1
-	t.allot(left, oldPrefixIndex, newPrefixIndex, val)
+	t.allot(left, old, new)
 	right := left + 1
-	t.allot(right, oldPrefixIndex, newPrefixIndex, val)
+	t.allot(right, old, new)
 }
 
 // insert adds the route addr/prefixLen to t, with value val.
-func (t *strideTable[T]) insert(addr uint8, prefixLen int, val *T) {
+func (t *strideTable[T]) insert(addr uint8, prefixLen int, val T) {
 	idx := prefixIndex(addr, prefixLen)
-	old := t.entries[idx].value
-	oldIdx := t.entries[idx].prefixIndex
-	if oldIdx == idx && old == val {
-		// This exact prefix+value is already in the table.
-		return
-	}
-	t.allot(idx, oldIdx, idx, val)
-	if oldIdx != idx {
-		// This route entry was freshly created (not just updated), that's a new
-		// reference.
+	if !t.hasPrefixRootedAt(idx) {
+		// This route entry is being freshly created (not just
+		// updated), that's a new reference.
 		t.routeRefs++
 	}
+
+	old := t.entries[idx]
+
+	// For allot to work correctly, each distinct prefix in the
+	// strideTable must have a different value pointer, even if val is
+	// identical. This new()+assignment guarantees that each inserted
+	// prefix gets a unique address.
+	p := new(T)
+	*p = val
+
+	t.allot(idx, old, p)
 	return
 }
 
-// delete removes the route addr/prefixLen from t. Returns the value
-// that was associated with the deleted prefix, or nil if the prefix
-// wasn't in the strideTable.
-func (t *strideTable[T]) delete(addr uint8, prefixLen int) *T {
+// delete removes the route addr/prefixLen from t. Reports whether the
+// prefix existed in the table prior to deletion.
+func (t *strideTable[T]) delete(addr uint8, prefixLen int) (wasPresent bool) {
 	idx := prefixIndex(addr, prefixLen)
-	recordedIdx := t.entries[idx].prefixIndex
-	if recordedIdx != idx {
+	if !t.hasPrefixRootedAt(idx) {
 		// Route entry doesn't exist
-		return nil
+		return false
 	}
-	val := t.entries[idx].value
 
-	parentIdx := idx >> 1
-	t.allot(idx, idx, t.entries[parentIdx].prefixIndex, t.entries[parentIdx].value)
+	val := t.entries[idx]
+	var parentVal *T
+	if parentIdx := parentIndex(idx); parentIdx != 0 {
+		parentVal = t.entries[parentIdx]
+	}
+
+	t.allot(idx, val, parentVal)
 	t.routeRefs--
-	return val
+	return true
 }
 
-// get does a route lookup for addr and returns the associated value, or nil if
-// no route matched.
-func (t *strideTable[T]) get(addr uint8) *T {
-	return t.entries[hostIndex(addr)].value
+// get does a route lookup for addr and (value, true) if a matching
+// route exists, or (zero, false) otherwise.
+func (t *strideTable[T]) get(addr uint8) (ret T, ok bool) {
+	if val := t.entries[hostIndex(addr)]; val != nil {
+		return *val, true
+	}
+	return ret, false
+}
+
+// getValAndChild returns both the prefix value and child strideTable
+// for addr. valOK reports whether a prefix value exists for addr, and
+// child is non-nil if a child exists for addr.
+func (t *strideTable[T]) getValAndChild(addr uint8) (val T, valOK bool, child *strideTable[T]) {
+	vp := t.entries[hostIndex(addr)]
+	if vp != nil {
+		val = *vp
+		valOK = true
+	}
+	child = t.children[addr]
+	return
 }
 
 // TableDebugString returns the contents of t, formatted as a table with one
@@ -208,10 +232,10 @@ func (t *strideTable[T]) tableDebugString() string {
 			continue
 		}
 		v := "(nil)"
-		if ent.value != nil {
-			v = fmt.Sprint(*ent.value)
+		if ent != nil {
+			v = fmt.Sprint(*ent)
 		}
-		fmt.Fprintf(&ret, "idx=%3d (%s), parent=%3d (%s), val=%v\n", i, formatPrefixTable(inversePrefixIndex(i)), ent.prefixIndex, formatPrefixTable(inversePrefixIndex((ent.prefixIndex))), v)
+		fmt.Fprintf(&ret, "idx=%3d (%s), val=%v\n", i, formatPrefixTable(inversePrefixIndex(i)), v)
 	}
 	return ret.String()
 }
@@ -227,8 +251,8 @@ func (t *strideTable[T]) treeDebugString() string {
 
 func (t *strideTable[T]) treeDebugStringRec(w io.Writer, idx, indent int) {
 	addr, len := inversePrefixIndex(idx)
-	if t.entries[idx].prefixIndex != 0 && t.entries[idx].prefixIndex == idx {
-		fmt.Fprintf(w, "%s%d/%d (%02x/%d) = %v\n", strings.Repeat(" ", indent), addr, len, addr, len, *t.entries[idx].value)
+	if t.hasPrefixRootedAt(idx) {
+		fmt.Fprintf(w, "%s%d/%d (%02x/%d) = %v\n", strings.Repeat(" ", indent), addr, len, addr, len, *t.entries[idx])
 		indent += 2
 	}
 	if idx >= firstHostIndex {
@@ -251,6 +275,12 @@ func prefixIndex(addr uint8, prefixLen int) int {
 	return (int(addr) >> (8 - prefixLen)) + (1 << prefixLen)
 }
 
+// parentIndex returns the index of idx's parent prefix, or 0 if idx
+// is the index of 0/0.
+func parentIndex(idx int) int {
+	return idx >> 1
+}
+
 // hostIndex returns the array index of the host route for addr.
 // It is equivalent to prefixIndex(addr, 8).
 func hostIndex(addr uint8) int {