feat: implement ClosestN

Returns the closest N peers from the table. This allows one to use this trie as
an routing table.

kademlia/bucket.go

@@ -19,3 +19,37 @@ func BucketAtDepth(node key.Key, table *trie.Trie, depth int) *trie.Trie {
 		}
 	}
 }
+
+// ClosestN will return the count closest keys to the given key.
+func ClosestN(node key.Key, table *trie.Trie, count int) []key.Key {
+	return closestAtDepth(node, table, 0, count, make([]key.Key, 0, count))
+}
+
+func closestAtDepth(node key.Key, table *trie.Trie, depth int, count int, found []key.Key) []key.Key {
+	// If we've already found enough peers, abort.
+	if count == len(found) {
+		return found
+	}
+
+	// Find the closest direction.
+	dir := node.BitAt(depth)
+	var chosenDir byte
+	if table.Branch[dir] != nil {
+		// There are peers in the "closer" direction.
+		chosenDir = dir
+	} else if table.Branch[1-dir] != nil {
+		// There are peers in the "less closer" direction.
+		chosenDir = 1 - dir
+	} else if table.Key != nil {
+		// We've found a leaf
+		return append(found, table.Key)
+	} else {
+		// We've found an empty node?
+		return found
+	}
+
+	// Add peers from the closest direction first, then from the other direction.
+	found = closestAtDepth(node, table.Branch[chosenDir], depth+1, count, found)
+	found = closestAtDepth(node, table.Branch[1-chosenDir], depth+1, count, found)
+	return found
+}

kademlia/bucket_test.go

+package kademlia
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/libp2p/go-libp2p-xor/key"
+	"github.com/libp2p/go-libp2p-xor/trie"
+)
+
+func randomTrie(count int, keySizeByte int) *trie.Trie {
+	t := trie.New()
+	for i := 0; i < count; i++ {
+		t.Add(randomKey(keySizeByte))
+	}
+	return t
+}
+
+func TestClosestN(t *testing.T) {
+	keySizeByte := 16
+	root := randomTrie(100, keySizeByte)
+	all := root.List()
+	for count := 0; count <= 100; count += 10 {
+		target := randomKey(keySizeByte)
+		closest := ClosestN(target, root, count)
+		if len(closest) != count {
+			t.Fatalf("expected %d closest, found %d", count, len(closest))
+		}
+		sort.Slice(all, func(i, j int) bool {
+			return string(key.Xor(all[i], target)) < string(key.Xor(all[j], target))
+		})
+		for i := range closest {
+			if !key.Equal(closest[i], all[i]) {
+				t.Errorf("wrong closest peer at offset %d: got %s, expected %s", i,
+					closest[i], all[i])
+			}
+		}
+	}
+}
+
+func closestTrivial(target key.Key, keys []key.Key, count int) []key.Key {
+	sort.Slice(keys, func(i, j int) bool {
+		return string(key.Xor(keys[i], target)) < string(key.Xor(keys[j], target))
+	})
+	return keys[:count]
+}
+
+// ensure the benchmark works.
+var _x int
+
+func BenchmarkClosestN(b *testing.B) {
+	keySizeByte := 16
+	root := randomTrie(100000, keySizeByte)
+	count := 20
+	target := randomKey(keySizeByte)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_x += len(ClosestN(target, root, count))
+	}
+}
+
+func BenchmarkClosestTrivial(b *testing.B) {
+	keySizeByte := 16
+	root := randomTrie(100000, keySizeByte)
+	keys := root.List()
+	count := 20
+	target := randomKey(keySizeByte)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_x += len(closestTrivial(target, keys, count))
+	}
+}