refactor(provider): closestPeerToPrefix coverage trie

provider/internal/keyspace/trie.go

@@ -1,6 +1,8 @@
 package keyspace
 
 import (
+	"slices"
+
 	"github.com/libp2p/go-libp2p/core/peer"
 	mh "github.com/multiformats/go-multihash"
 
@@ -185,46 +187,92 @@ func pruneSubtrieAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0,
 }
 
 // TrieGaps returns all prefixes that aren't covered by a key (prefix) in the
-// trie. Combining the prefixes included in the trie with the gap prefixes
-// results in a full keyspace coverage.
+// trie, at the `target` location. Combining the prefixes included in the trie
+// with the gap prefixes results in a full keyspace coverage of the `target`
+// subtrie.
+//
+// Results are sorted according to the provided `order`.
 //
-// E.g Trie: ["00", "100"], GapsInTrie: ["01", "101", "11"]
-func TrieGaps[D any](t *trie.Trie[bitstr.Key, D]) []bitstr.Key {
+// Example:
+//   - Trie: ["0000", "0010", "100"]
+//   - Target: "0"
+//   - Order: "0000",
+//   - GapsInTrie: ["0001", "0011", "01"]
+func TrieGaps[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], target bitstr.Key, order K) []bitstr.Key {
 	if t.IsLeaf() {
-		if t.HasKey() {
-			return SiblingPrefixes(*t.Key())
+		if k := t.Key(); k != nil {
+			if IsBitstrPrefix(target, *k) {
+				siblingPrefixes := SiblingPrefixes(*k)[len(target):]
+				sortBitstrKeysByOrder(siblingPrefixes, order)
+				return siblingPrefixes
+			}
+			if IsBitstrPrefix(*k, target) {
+				// The only key in the trie is a prefix of target, meaning the whole
+				// target is covered.
+				return nil
+			}
 		}
-		return []bitstr.Key{""}
+		return []bitstr.Key{target}
 	}
-	return trieGapsAtDepth(t, 0)
+	return trieGapsAtDepth(t, 0, target, order)
 }
 
-func trieGapsAtDepth[D any](t *trie.Trie[bitstr.Key, D], depth int) []bitstr.Key {
+func trieGapsAtDepth[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], depth int, target bitstr.Key, order K) []bitstr.Key {
 	var gaps []bitstr.Key
-	for i := range 2 {
+	insideTarget := depth >= target.BitLen()
+	b := int(order.Bit(depth))
+	for _, i := range []int{b, 1 - b} {
+		if !insideTarget && i != int(target.Bit(depth)) {
+			continue
+		}
 		bstr := bitstr.Key(byte('0' + i))
 		if b := t.Branch(i); b == nil {
 			gaps = append(gaps, bstr)
 		} else if b.IsLeaf() {
 			if b.HasKey() {
 				k := *b.Key()
 				if len(k) > depth+1 {
-					for _, siblingPrefix := range SiblingPrefixes(k)[depth+1:] {
+					siblingPrefixes := SiblingPrefixes(k)[depth+1:]
+					sortBitstrKeysByOrder(siblingPrefixes, order)
+					for _, siblingPrefix := range siblingPrefixes {
 						gaps = append(gaps, siblingPrefix[depth:])
 					}
 				}
 			} else {
 				gaps = append(gaps, bstr)
 			}
 		} else {
-			for _, gap := range trieGapsAtDepth(b, depth+1) {
+			for _, gap := range trieGapsAtDepth(b, depth+1, target, order) {
 				gaps = append(gaps, bstr+gap)
 			}
 		}
 	}
 	return gaps
 }
 
+// sortBitstrKeysByOrder sorts the provided bitstr keys according to the
+// provided order.
+func sortBitstrKeysByOrder[K kad.Key[K]](keys []bitstr.Key, order K) {
+	slices.SortFunc(keys, func(a, b bitstr.Key) int {
+		maxLen := min(len(a), len(b), order.BitLen())
+		for i := range maxLen {
+			if a[i] != b[i] {
+				if a.Bit(i) == order.Bit(i) {
+					return -1
+				}
+				return 1
+			}
+		}
+		if len(a) == len(b) || maxLen == order.BitLen() {
+			return 0
+		}
+		if len(a) < len(b) {
+			return 1
+		}
+		return -1
+	})
+}
+
 // mapMerge merges all key-value pairs from the source map into the destination
 // map. Values from the source are appended to existing slices in the
 // destination.

provider/internal/keyspace/trie_test.go

@@ -444,15 +444,15 @@ func TestTrieGaps(t *testing.T) {
 	t.Run("Gap in empty trie", func(t *testing.T) {
 		keys := []bitstr.Key{}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{""}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{""}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("No gaps in flat trie", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"0",
 			"1",
 		}
 		tr := initTrie(keys)
-		require.Empty(t, TrieGaps(tr))
+		require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("No gaps in unbalanced trie", func(t *testing.T) {
 		keys := []bitstr.Key{
@@ -462,35 +462,37 @@ func TestTrieGaps(t *testing.T) {
 			"111",
 		}
 		tr := initTrie(keys)
-		require.Empty(t, TrieGaps(tr))
+		require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("No gaps in trie with empty key", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"",
 		}
 		tr := initTrie(keys)
-		require.Empty(t, TrieGaps(tr))
+		require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap with single key - 0", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"0",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"1"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"1"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap with single key - 1", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"1",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"0"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"0"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap with single key - 11101101", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"11101101",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, SiblingPrefixes(keys[0]), TrieGaps(tr))
+		siblingPrefixes := SiblingPrefixes(keys[0])
+		sortBitstrKeysByOrder(siblingPrefixes, bit256.ZeroKey())
+		require.Equal(t, siblingPrefixes, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap missing single key - 0", func(t *testing.T) {
 		keys := []bitstr.Key{
@@ -499,7 +501,7 @@ func TestTrieGaps(t *testing.T) {
 			"110",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap missing single key - 1", func(t *testing.T) {
 		keys := []bitstr.Key{
@@ -509,7 +511,7 @@ func TestTrieGaps(t *testing.T) {
 			"1101",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap missing single key - 2", func(t *testing.T) {
 		keys := []bitstr.Key{
@@ -521,7 +523,7 @@ func TestTrieGaps(t *testing.T) {
 			"1101",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap missing multiple keys - 0", func(t *testing.T) {
 		keys := []bitstr.Key{
@@ -532,15 +534,133 @@ func TestTrieGaps(t *testing.T) {
 			"1101",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"001", "111"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"001", "111"}, TrieGaps(tr, "", bit256.ZeroKey()))
 	})
 	t.Run("Gap missing multiple keys - 1", func(t *testing.T) {
 		keys := []bitstr.Key{
 			"000",
 			"1101",
 		}
 		tr := initTrie(keys)
-		require.Equal(t, []bitstr.Key{"01", "001", "10", "111", "1100"}, TrieGaps(tr))
+		require.Equal(t, []bitstr.Key{"001", "01", "10", "1100", "111"}, TrieGaps(tr, "", bit256.ZeroKey()))
+	})
+	t.Run("Gap missing multiple keys - 2", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"0000",
+			"1000",
+		}
+		tr := initTrie(keys)
+		require.Equal(t, []bitstr.Key{"0001", "001", "01", "1001", "101", "11"}, TrieGaps(tr, "", bit256.ZeroKey()))
+	})
+
+	t.Run("Single key inside target", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"0000",
+		}
+		tr := initTrie(keys)
+		require.Equal(t, []bitstr.Key{"0001", "001"}, TrieGaps(tr, "00", bit256.ZeroKey()))
+	})
+	t.Run("Single key outside target", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"0000",
+		}
+		tr := initTrie(keys)
+		require.Equal(t, []bitstr.Key{"11"}, TrieGaps(tr, "11", bit256.ZeroKey()))
+	})
+
+	t.Run("Target subset", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"0000",
+			"0010",
+			"100",
+			"11",
+		}
+		tr := initTrie(keys)
+		require.Equal(t, []bitstr.Key{"0001", "0011", "01"}, TrieGaps(tr, "0", bit256.ZeroKey()))
+	})
+
+	t.Run("Target subset reverse order", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"0000",
+			"0001",
+			"1000",
+		}
+		tr := initTrie(keys)
+		require.Equal(t, []bitstr.Key{"01", "001"}, TrieGaps(tr, "0", bitstr.Key("1111")))
+	})
+
+	t.Run("Target longer than only key in trie", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"00",
+		}
+		tr := initTrie(keys)
+		require.Empty(t, TrieGaps(tr, "000", bit256.ZeroKey()))
+	})
+
+	t.Run("Target is superstring of key in trie", func(t *testing.T) {
+		keys := []bitstr.Key{
+			"00",
+			"01",
+			"101",
+		}
+		tr := initTrie(keys)
+		require.Empty(t, TrieGaps(tr, "000", bit256.ZeroKey()))
+	})
+}
+
+func TestSortBitstrKeysByOrder(t *testing.T) {
+	t.Run("Empty", func(t *testing.T) {
+		keys := []bitstr.Key{}
+		sortBitstrKeysByOrder(keys, bit256.ZeroKey())
+		require.Equal(t, []bitstr.Key{}, keys)
+	})
+
+	t.Run("Single key", func(t *testing.T) {
+		keys := []bitstr.Key{"0"}
+		sortBitstrKeysByOrder(keys, bit256.ZeroKey())
+		require.Equal(t, []bitstr.Key{"0"}, keys)
+	})
+
+	t.Run("Two sorted keys", func(t *testing.T) {
+		keys := []bitstr.Key{"0", "1"}
+		sortBitstrKeysByOrder(keys, bit256.ZeroKey())
+		require.Equal(t, []bitstr.Key{"0", "1"}, keys)
+	})
+
+	t.Run("Two keys - zero order", func(t *testing.T) {
+		keys := []bitstr.Key{"1", "0"}
+		sortBitstrKeysByOrder(keys, bit256.ZeroKey())
+		require.Equal(t, []bitstr.Key{"0", "1"}, keys)
+	})
+
+	t.Run("Two keys - reverse order", func(t *testing.T) {
+		keys := []bitstr.Key{"0", "1"}
+		sortBitstrKeysByOrder(keys, bitstr.Key("1"))
+		require.Equal(t, []bitstr.Key{"1", "0"}, keys)
+	})
+
+	t.Run("Different key lengths", func(t *testing.T) {
+		keys := []bitstr.Key{"00", "000", "0"}
+		sortBitstrKeysByOrder(keys, bitstr.Key("0000"))
+		require.Equal(t, []bitstr.Key{"000", "00", "0"}, keys)
+	})
+
+	t.Run("Different key lengths", func(t *testing.T) {
+		keys := []bitstr.Key{"01", "1"}
+		sortBitstrKeysByOrder(keys, bitstr.Key("11"))
+		require.Equal(t, []bitstr.Key{"1", "01"}, keys)
+	})
+
+	t.Run("Short order", func(t *testing.T) {
+		keys := []bitstr.Key{"111", "110", "0"}
+		sortBitstrKeysByOrder(keys, bitstr.Key("00"))
+		require.Equal(t, []bitstr.Key{"0", "111", "110"}, keys)
+	})
+
+	t.Run("Identical keys", func(t *testing.T) {
+		keys := []bitstr.Key{"0", "0"}
+		sortBitstrKeysByOrder(keys, bitstr.Key("00"))
+		require.Equal(t, []bitstr.Key{"0", "0"}, keys)
 	})
 }