Skip to content

Commit 0114a72

Browse files
guillaumemichellidel
guillaumemichel
and
lidel
authored
refactor(provider): closestPeerToPrefix coverage trie (#1156)
* refactor(provider): closestPeerToPrefix coverage trie * fix infinite loop * fix: don't insert if trie contains prefix * fix: return no gap if target is covered by a shorter key * test: add closestPeersToPrefix error handling tests * fix: use debug log level for region exploration metrics --------- Co-authored-by: Marcin Rataj <[email protected]>
1 parent 548e361 commit 0114a72

File tree

4 files changed

+288
-69
lines changed

4 files changed

+288
-69
lines changed

‎provider/internal/keyspace/trie.go‎

Lines changed: 60 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package keyspace
22

33
import (
4+
"slices"
5+
46
"github.com/libp2p/go-libp2p/core/peer"
57
mh "github.com/multiformats/go-multihash"
68

@@ -185,46 +187,92 @@ func pruneSubtrieAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0,
185187
}
186188

187189
// TrieGaps returns all prefixes that aren't covered by a key (prefix) in the
188-
// trie. Combining the prefixes included in the trie with the gap prefixes
189-
// results in a full keyspace coverage.
190+
// trie, at the `target` location. Combining the prefixes included in the trie
191+
// with the gap prefixes results in a full keyspace coverage of the `target`
192+
// subtrie.
193+
//
194+
// Results are sorted according to the provided `order`.
190195
//
191-
// E.g Trie: ["00", "100"], GapsInTrie: ["01", "101", "11"]
192-
func TrieGaps[D any](t *trie.Trie[bitstr.Key, D]) []bitstr.Key {
196+
// Example:
197+
// - Trie: ["0000", "0010", "100"]
198+
// - Target: "0"
199+
// - Order: "0000",
200+
// - GapsInTrie: ["0001", "0011", "01"]
201+
func TrieGaps[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], target bitstr.Key, order K) []bitstr.Key {
193202
if t.IsLeaf() {
194-
if t.HasKey() {
195-
return SiblingPrefixes(*t.Key())
203+
if k := t.Key(); k != nil {
204+
if IsBitstrPrefix(target, *k) {
205+
siblingPrefixes := SiblingPrefixes(*k)[len(target):]
206+
sortBitstrKeysByOrder(siblingPrefixes, order)
207+
return siblingPrefixes
208+
}
209+
if IsBitstrPrefix(*k, target) {
210+
// The only key in the trie is a prefix of target, meaning the whole
211+
// target is covered.
212+
return nil
213+
}
196214
}
197-
return []bitstr.Key{""}
215+
return []bitstr.Key{target}
198216
}
199-
return trieGapsAtDepth(t, 0)
217+
return trieGapsAtDepth(t, 0, target, order)
200218
}
201219

202-
func trieGapsAtDepth[D any](t *trie.Trie[bitstr.Key, D], depth int) []bitstr.Key {
220+
func trieGapsAtDepth[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], depth int, target bitstr.Key, order K) []bitstr.Key {
203221
var gaps []bitstr.Key
204-
for i := range 2 {
222+
insideTarget := depth >= target.BitLen()
223+
b := int(order.Bit(depth))
224+
for _, i := range []int{b, 1 - b} {
225+
if !insideTarget && i != int(target.Bit(depth)) {
226+
continue
227+
}
205228
bstr := bitstr.Key(byte('0' + i))
206229
if b := t.Branch(i); b == nil {
207230
gaps = append(gaps, bstr)
208231
} else if b.IsLeaf() {
209232
if b.HasKey() {
210233
k := *b.Key()
211234
if len(k) > depth+1 {
212-
for _, siblingPrefix := range SiblingPrefixes(k)[depth+1:] {
235+
siblingPrefixes := SiblingPrefixes(k)[depth+1:]
236+
sortBitstrKeysByOrder(siblingPrefixes, order)
237+
for _, siblingPrefix := range siblingPrefixes {
213238
gaps = append(gaps, siblingPrefix[depth:])
214239
}
215240
}
216241
} else {
217242
gaps = append(gaps, bstr)
218243
}
219244
} else {
220-
for _, gap := range trieGapsAtDepth(b, depth+1) {
245+
for _, gap := range trieGapsAtDepth(b, depth+1, target, order) {
221246
gaps = append(gaps, bstr+gap)
222247
}
223248
}
224249
}
225250
return gaps
226251
}
227252

253+
// sortBitstrKeysByOrder sorts the provided bitstr keys according to the
254+
// provided order.
255+
func sortBitstrKeysByOrder[K kad.Key[K]](keys []bitstr.Key, order K) {
256+
slices.SortFunc(keys, func(a, b bitstr.Key) int {
257+
maxLen := min(len(a), len(b), order.BitLen())
258+
for i := range maxLen {
259+
if a[i] != b[i] {
260+
if a.Bit(i) == order.Bit(i) {
261+
return -1
262+
}
263+
return 1
264+
}
265+
}
266+
if len(a) == len(b) || maxLen == order.BitLen() {
267+
return 0
268+
}
269+
if len(a) < len(b) {
270+
return 1
271+
}
272+
return -1
273+
})
274+
}
275+
228276
// mapMerge merges all key-value pairs from the source map into the destination
229277
// map. Values from the source are appended to existing slices in the
230278
// destination.

‎provider/internal/keyspace/trie_test.go‎

Lines changed: 132 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -444,15 +444,15 @@ func TestTrieGaps(t *testing.T) {
444444
t.Run("Gap in empty trie", func(t *testing.T) {
445445
keys := []bitstr.Key{}
446446
tr := initTrie(keys)
447-
require.Equal(t, []bitstr.Key{""}, TrieGaps(tr))
447+
require.Equal(t, []bitstr.Key{""}, TrieGaps(tr, "", bit256.ZeroKey()))
448448
})
449449
t.Run("No gaps in flat trie", func(t *testing.T) {
450450
keys := []bitstr.Key{
451451
"0",
452452
"1",
453453
}
454454
tr := initTrie(keys)
455-
require.Empty(t, TrieGaps(tr))
455+
require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
456456
})
457457
t.Run("No gaps in unbalanced trie", func(t *testing.T) {
458458
keys := []bitstr.Key{
@@ -462,35 +462,37 @@ func TestTrieGaps(t *testing.T) {
462462
"111",
463463
}
464464
tr := initTrie(keys)
465-
require.Empty(t, TrieGaps(tr))
465+
require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
466466
})
467467
t.Run("No gaps in trie with empty key", func(t *testing.T) {
468468
keys := []bitstr.Key{
469469
"",
470470
}
471471
tr := initTrie(keys)
472-
require.Empty(t, TrieGaps(tr))
472+
require.Empty(t, TrieGaps(tr, "", bit256.ZeroKey()))
473473
})
474474
t.Run("Gap with single key - 0", func(t *testing.T) {
475475
keys := []bitstr.Key{
476476
"0",
477477
}
478478
tr := initTrie(keys)
479-
require.Equal(t, []bitstr.Key{"1"}, TrieGaps(tr))
479+
require.Equal(t, []bitstr.Key{"1"}, TrieGaps(tr, "", bit256.ZeroKey()))
480480
})
481481
t.Run("Gap with single key - 1", func(t *testing.T) {
482482
keys := []bitstr.Key{
483483
"1",
484484
}
485485
tr := initTrie(keys)
486-
require.Equal(t, []bitstr.Key{"0"}, TrieGaps(tr))
486+
require.Equal(t, []bitstr.Key{"0"}, TrieGaps(tr, "", bit256.ZeroKey()))
487487
})
488488
t.Run("Gap with single key - 11101101", func(t *testing.T) {
489489
keys := []bitstr.Key{
490490
"11101101",
491491
}
492492
tr := initTrie(keys)
493-
require.Equal(t, SiblingPrefixes(keys[0]), TrieGaps(tr))
493+
siblingPrefixes := SiblingPrefixes(keys[0])
494+
sortBitstrKeysByOrder(siblingPrefixes, bit256.ZeroKey())
495+
require.Equal(t, siblingPrefixes, TrieGaps(tr, "", bit256.ZeroKey()))
494496
})
495497
t.Run("Gap missing single key - 0", func(t *testing.T) {
496498
keys := []bitstr.Key{
@@ -499,7 +501,7 @@ func TestTrieGaps(t *testing.T) {
499501
"110",
500502
}
501503
tr := initTrie(keys)
502-
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
504+
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
503505
})
504506
t.Run("Gap missing single key - 1", func(t *testing.T) {
505507
keys := []bitstr.Key{
@@ -509,7 +511,7 @@ func TestTrieGaps(t *testing.T) {
509511
"1101",
510512
}
511513
tr := initTrie(keys)
512-
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
514+
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
513515
})
514516
t.Run("Gap missing single key - 2", func(t *testing.T) {
515517
keys := []bitstr.Key{
@@ -521,7 +523,7 @@ func TestTrieGaps(t *testing.T) {
521523
"1101",
522524
}
523525
tr := initTrie(keys)
524-
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr))
526+
require.Equal(t, []bitstr.Key{"111"}, TrieGaps(tr, "", bit256.ZeroKey()))
525527
})
526528
t.Run("Gap missing multiple keys - 0", func(t *testing.T) {
527529
keys := []bitstr.Key{
@@ -532,15 +534,133 @@ func TestTrieGaps(t *testing.T) {
532534
"1101",
533535
}
534536
tr := initTrie(keys)
535-
require.Equal(t, []bitstr.Key{"001", "111"}, TrieGaps(tr))
537+
require.Equal(t, []bitstr.Key{"001", "111"}, TrieGaps(tr, "", bit256.ZeroKey()))
536538
})
537539
t.Run("Gap missing multiple keys - 1", func(t *testing.T) {
538540
keys := []bitstr.Key{
539541
"000",
540542
"1101",
541543
}
542544
tr := initTrie(keys)
543-
require.Equal(t, []bitstr.Key{"01", "001", "10", "111", "1100"}, TrieGaps(tr))
545+
require.Equal(t, []bitstr.Key{"001", "01", "10", "1100", "111"}, TrieGaps(tr, "", bit256.ZeroKey()))
546+
})
547+
t.Run("Gap missing multiple keys - 2", func(t *testing.T) {
548+
keys := []bitstr.Key{
549+
"0000",
550+
"1000",
551+
}
552+
tr := initTrie(keys)
553+
require.Equal(t, []bitstr.Key{"0001", "001", "01", "1001", "101", "11"}, TrieGaps(tr, "", bit256.ZeroKey()))
554+
})
555+
556+
t.Run("Single key inside target", func(t *testing.T) {
557+
keys := []bitstr.Key{
558+
"0000",
559+
}
560+
tr := initTrie(keys)
561+
require.Equal(t, []bitstr.Key{"0001", "001"}, TrieGaps(tr, "00", bit256.ZeroKey()))
562+
})
563+
t.Run("Single key outside target", func(t *testing.T) {
564+
keys := []bitstr.Key{
565+
"0000",
566+
}
567+
tr := initTrie(keys)
568+
require.Equal(t, []bitstr.Key{"11"}, TrieGaps(tr, "11", bit256.ZeroKey()))
569+
})
570+
571+
t.Run("Target subset", func(t *testing.T) {
572+
keys := []bitstr.Key{
573+
"0000",
574+
"0010",
575+
"100",
576+
"11",
577+
}
578+
tr := initTrie(keys)
579+
require.Equal(t, []bitstr.Key{"0001", "0011", "01"}, TrieGaps(tr, "0", bit256.ZeroKey()))
580+
})
581+
582+
t.Run("Target subset reverse order", func(t *testing.T) {
583+
keys := []bitstr.Key{
584+
"0000",
585+
"0001",
586+
"1000",
587+
}
588+
tr := initTrie(keys)
589+
require.Equal(t, []bitstr.Key{"01", "001"}, TrieGaps(tr, "0", bitstr.Key("1111")))
590+
})
591+
592+
t.Run("Target longer than only key in trie", func(t *testing.T) {
593+
keys := []bitstr.Key{
594+
"00",
595+
}
596+
tr := initTrie(keys)
597+
require.Empty(t, TrieGaps(tr, "000", bit256.ZeroKey()))
598+
})
599+
600+
t.Run("Target is superstring of key in trie", func(t *testing.T) {
601+
keys := []bitstr.Key{
602+
"00",
603+
"01",
604+
"101",
605+
}
606+
tr := initTrie(keys)
607+
require.Empty(t, TrieGaps(tr, "000", bit256.ZeroKey()))
608+
})
609+
}
610+
611+
func TestSortBitstrKeysByOrder(t *testing.T) {
612+
t.Run("Empty", func(t *testing.T) {
613+
keys := []bitstr.Key{}
614+
sortBitstrKeysByOrder(keys, bit256.ZeroKey())
615+
require.Equal(t, []bitstr.Key{}, keys)
616+
})
617+
618+
t.Run("Single key", func(t *testing.T) {
619+
keys := []bitstr.Key{"0"}
620+
sortBitstrKeysByOrder(keys, bit256.ZeroKey())
621+
require.Equal(t, []bitstr.Key{"0"}, keys)
622+
})
623+
624+
t.Run("Two sorted keys", func(t *testing.T) {
625+
keys := []bitstr.Key{"0", "1"}
626+
sortBitstrKeysByOrder(keys, bit256.ZeroKey())
627+
require.Equal(t, []bitstr.Key{"0", "1"}, keys)
628+
})
629+
630+
t.Run("Two keys - zero order", func(t *testing.T) {
631+
keys := []bitstr.Key{"1", "0"}
632+
sortBitstrKeysByOrder(keys, bit256.ZeroKey())
633+
require.Equal(t, []bitstr.Key{"0", "1"}, keys)
634+
})
635+
636+
t.Run("Two keys - reverse order", func(t *testing.T) {
637+
keys := []bitstr.Key{"0", "1"}
638+
sortBitstrKeysByOrder(keys, bitstr.Key("1"))
639+
require.Equal(t, []bitstr.Key{"1", "0"}, keys)
640+
})
641+
642+
t.Run("Different key lengths", func(t *testing.T) {
643+
keys := []bitstr.Key{"00", "000", "0"}
644+
sortBitstrKeysByOrder(keys, bitstr.Key("0000"))
645+
require.Equal(t, []bitstr.Key{"000", "00", "0"}, keys)
646+
})
647+
648+
t.Run("Different key lengths", func(t *testing.T) {
649+
keys := []bitstr.Key{"01", "1"}
650+
sortBitstrKeysByOrder(keys, bitstr.Key("11"))
651+
require.Equal(t, []bitstr.Key{"1", "01"}, keys)
652+
})
653+
654+
t.Run("Short order", func(t *testing.T) {
655+
keys := []bitstr.Key{"111", "110", "0"}
656+
sortBitstrKeysByOrder(keys, bitstr.Key("00"))
657+
require.Equal(t, []bitstr.Key{"0", "111", "110"}, keys)
658+
})
659+
660+
t.Run("Identical keys", func(t *testing.T) {
661+
keys := []bitstr.Key{"0", "0"}
662+
sortBitstrKeysByOrder(keys, bitstr.Key("00"))
663+
require.Equal(t, []bitstr.Key{"0", "0"}, keys)
544664
})
545665
}
546666

0 commit comments

Comments
 (0)