provider: AllEntries and AllValues helpers

Author: guillaumemichel

provider/internal/helpers/trie.go

@@ -0,0 +1,35 @@
+package helpers
+
+import (
+	"github.com/probe-lab/go-libdht/kad"
+	"github.com/probe-lab/go-libdht/kad/trie"
+)
+
+// AllEntries returns all entries (key + value) stored in the trie `t` sorted
+// by their keys in the supplied `order`.
+func AllEntries[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0, D], order K1) []*trie.Entry[K0, D] {
+	return allEntriesAtDepth(t, order, 0)
+}
+
+func allEntriesAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0, D], order K1, depth int) []*trie.Entry[K0, D] {
+	if t == nil || t.IsEmptyLeaf() {
+		return nil
+	}
+	if t.IsNonEmptyLeaf() {
+		return []*trie.Entry[K0, D]{{Key: *t.Key(), Data: t.Data()}}
+	}
+	b := int(order.Bit(depth))
+	return append(allEntriesAtDepth(t.Branch(b), order, depth+1),
+		allEntriesAtDepth(t.Branch(1-b), order, depth+1)...)
+}
+
+// AllValues returns all values stored in the trie `t` sorted by their keys in
+// the supplied `order`.
+func AllValues[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0, D], order K1) []D {
+	entries := AllEntries(t, order)
+	out := make([]D, len(entries))
+	for i, entry := range entries {
+		out[i] = entry.Data
+	}
+	return out
+}

provider/internal/helpers/trie_test.go

@@ -0,0 +1,55 @@
+package helpers
+
+import (
+	"testing"
+
+	"github.com/probe-lab/go-libdht/kad/key/bitstr"
+	"github.com/probe-lab/go-libdht/kad/trie"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestAllEntries(t *testing.T) {
+	tr := trie.New[bitstr.Key, string]()
+	elements := []struct {
+		key   bitstr.Key
+		fruit string
+	}{
+		{
+			key:   bitstr.Key("000"), // convert to Key if it’s a distinct type
+			fruit: "apple",
+		},
+		{
+			key:   bitstr.Key("010"),
+			fruit: "banana",
+		},
+		{
+			key:   bitstr.Key("101"),
+			fruit: "cherry",
+		},
+		{
+			key:   bitstr.Key("111"),
+			fruit: "durian",
+		},
+	}
+
+	for _, e := range elements {
+		tr.Add(e.key, e.fruit)
+	}
+
+	// Test in 0 -> 1 order
+	entries := AllEntries(tr, bitstr.Key("000"))
+	require.Equal(t, len(elements), len(entries))
+	for i := range entries {
+		require.Equal(t, entries[i].Key, elements[i].key)
+		require.Equal(t, entries[i].Data, elements[i].fruit)
+	}
+
+	// Test in reverse order (1 -> 0)
+	entries = AllEntries(tr, bitstr.Key("111"))
+	require.Equal(t, len(elements), len(entries))
+	for i := range entries {
+		require.Equal(t, entries[i].Key, elements[len(elements)-1-i].key)
+		require.Equal(t, entries[i].Data, elements[len(elements)-1-i].fruit)
+	}
+}