Skip to content

Commit d025c87

Browse files
infrmtcsEgeCaner
infrmtcs
authored and
EgeCaner
committed
feat: Typed prefixed bucket
1 parent ea2487e commit d025c87

File tree

5 files changed

+369
-0
lines changed

5 files changed

+369
-0
lines changed

‎db/typed/prefix/bucket.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package prefix
2+
3+
import (
4+
"iter"
5+
6+
"github.com/NethermindEth/juno/db"
7+
"github.com/NethermindEth/juno/db/typed"
8+
"github.com/NethermindEth/juno/db/typed/key"
9+
"github.com/NethermindEth/juno/db/typed/value"
10+
"github.com/NethermindEth/juno/utils"
11+
)
12+
13+
type Entry[V any] struct {
14+
Key []byte
15+
Value V
16+
}
17+
18+
type scanner[B any] interface {
19+
scan(database db.KeyValueReader, prefix []byte) iter.Seq2[Entry[B], error]
20+
}
21+
22+
type PrefixedBucket[T tag[V], K any, V any, KS key.Serializer[K], VS value.Serializer[V]] struct {
23+
typed.Bucket[K, V, KS, VS]
24+
}
25+
26+
func NewPrefixedBucket[T tag[V], K any, V any, KS key.Serializer[K], VS value.Serializer[V]](
27+
bucket typed.Bucket[K, V, KS, VS],
28+
prefixFactory proxy[V, T],
29+
) PrefixedBucket[T, K, V, KS, VS] {
30+
return PrefixedBucket[T, K, V, KS, VS]{
31+
Bucket: bucket,
32+
}
33+
}
34+
35+
func (b *PrefixedBucket[T, K, V, KS, VS]) Prefix() T {
36+
return T{
37+
scanState: scanState[V]{
38+
prefix: b.Bucket.Key(),
39+
scanner: b,
40+
},
41+
}
42+
}
43+
44+
func (b *PrefixedBucket[T, K, V, KS, VS]) scan( //nolint:unused // False positive, used via scanner
45+
database db.KeyValueReader,
46+
prefix []byte,
47+
) iter.Seq2[Entry[V], error] {
48+
return func(yield func(Entry[V], error) bool) {
49+
it, err := database.NewIterator(prefix, true)
50+
if err != nil {
51+
yield(Entry[V]{}, err)
52+
return
53+
}
54+
55+
isError := false
56+
defer func() {
57+
if !isError {
58+
if err := it.Close(); err != nil {
59+
yield(Entry[V]{}, err)
60+
}
61+
}
62+
}()
63+
64+
for it.First(); it.Valid(); it.Next() {
65+
var entry Entry[V]
66+
entry.Key = it.Key()
67+
68+
valueBytes, err := it.Value()
69+
if err != nil {
70+
yield(Entry[V]{}, utils.RunAndWrapOnError(it.Close, err))
71+
isError = true
72+
return
73+
}
74+
75+
if err := (VS{}).Unmarshal(valueBytes, &entry.Value); err != nil {
76+
yield(Entry[V]{}, utils.RunAndWrapOnError(it.Close, err))
77+
isError = true
78+
return
79+
}
80+
81+
if !yield(entry, nil) {
82+
return
83+
}
84+
}
85+
}
86+
}

‎db/typed/prefix/bucket_test.go‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
package prefix_test
2+
3+
import (
4+
"cmp"
5+
cryptorand "crypto/rand"
6+
"iter"
7+
"maps"
8+
"math/rand/v2"
9+
"slices"
10+
"testing"
11+
12+
"github.com/NethermindEth/juno/core/felt"
13+
"github.com/NethermindEth/juno/db"
14+
"github.com/NethermindEth/juno/db/memory"
15+
"github.com/NethermindEth/juno/db/typed"
16+
"github.com/NethermindEth/juno/db/typed/key"
17+
"github.com/NethermindEth/juno/db/typed/prefix"
18+
"github.com/NethermindEth/juno/db/typed/value"
19+
"github.com/stretchr/testify/require"
20+
)
21+
22+
const (
23+
layerKeyCount = 50
24+
testPerLayer = 5
25+
)
26+
27+
type testKey struct {
28+
number uint64
29+
hash felt.Felt
30+
slot []byte
31+
}
32+
33+
type testEntry struct {
34+
key testKey
35+
value felt.Felt
36+
}
37+
38+
func (k testKey) Marshal() []byte {
39+
return slices.Concat(
40+
key.Uint64.Marshal(k.number),
41+
key.Felt.Marshal(&k.hash),
42+
key.Bytes.Marshal(k.slot),
43+
)
44+
}
45+
46+
var bucket = prefix.NewPrefixedBucket(
47+
typed.NewBucket(
48+
db.Bucket(0),
49+
key.Marshal[testKey](),
50+
value.Felt,
51+
),
52+
prefix.Prefix(
53+
key.Uint64,
54+
prefix.Prefix(
55+
key.Felt,
56+
prefix.Prefix(
57+
key.Bytes,
58+
prefix.End[felt.Felt](),
59+
),
60+
),
61+
),
62+
)
63+
64+
func extractExpectedEntries(
65+
entries map[uint64]map[felt.Felt]map[string]testEntry,
66+
filterBlockNumber *uint64,
67+
filterHash *felt.Felt,
68+
filterSlot *string,
69+
) []testEntry {
70+
res := make([]testEntry, 0)
71+
for blockNumber := range filter(filterBlockNumber, cmp.Compare, entries) {
72+
for hash := range filter(filterHash, compareFelt, entries[blockNumber]) {
73+
for slot := range filter(filterSlot, cmp.Compare, entries[blockNumber][hash]) {
74+
res = append(res, entries[blockNumber][hash][slot])
75+
}
76+
}
77+
}
78+
return res
79+
}
80+
81+
func compareFelt(a, b felt.Felt) int {
82+
return a.Cmp(&b)
83+
}
84+
85+
func filter[K comparable, V any](filter *K, cmp func(K, K) int, entries map[K]V) iter.Seq[K] {
86+
if filter != nil {
87+
return func(yield func(K) bool) {
88+
yield(*filter)
89+
}
90+
}
91+
92+
return func(yield func(K) bool) {
93+
for _, key := range slices.SortedFunc(maps.Keys(entries), cmp) {
94+
if !yield(key) {
95+
return
96+
}
97+
}
98+
}
99+
}
100+
101+
func randomEntry[K comparable, V any](t *testing.T, entries map[K]V) (K, V) {
102+
t.Helper()
103+
for key, value := range entries {
104+
return key, value
105+
}
106+
require.FailNow(t, "no entries")
107+
return *new(K), *new(V)
108+
}
109+
110+
func validateResult(
111+
t *testing.T,
112+
expected []testEntry,
113+
actual iter.Seq2[prefix.Entry[felt.Felt], error],
114+
) {
115+
t.Helper()
116+
count := 0
117+
for entry, err := range actual {
118+
require.NoError(t, err)
119+
require.Greater(t, len(expected), count)
120+
require.Equal(t, bucket.Key(expected[count].key.Marshal()), entry.Key)
121+
require.Equal(t, expected[count].value, entry.Value)
122+
count++
123+
}
124+
require.Equal(t, len(expected), count)
125+
}
126+
127+
func TestPrefixedBucket(t *testing.T) {
128+
database := memory.New()
129+
content := make(map[uint64]map[felt.Felt]map[string]testEntry)
130+
131+
t.Run("Populate database", func(t *testing.T) {
132+
for range layerKeyCount {
133+
blockNumber := rand.Uint64()
134+
content[blockNumber] = make(map[felt.Felt]map[string]testEntry)
135+
136+
for range layerKeyCount {
137+
hash := felt.Random[felt.Felt]()
138+
content[blockNumber][hash] = make(map[string]testEntry)
139+
140+
for range layerKeyCount {
141+
slot := cryptorand.Text()
142+
value := felt.Random[felt.Felt]()
143+
144+
key := testKey{
145+
number: blockNumber,
146+
hash: hash,
147+
slot: []byte(slot),
148+
}
149+
entry := testEntry{
150+
key: key,
151+
value: value,
152+
}
153+
154+
content[blockNumber][hash][slot] = entry
155+
require.NoError(t, bucket.Put(database, key, &value))
156+
}
157+
}
158+
}
159+
})
160+
161+
t.Run("Full scan", func(t *testing.T) {
162+
validateResult(
163+
t,
164+
extractExpectedEntries(content, nil, nil, nil),
165+
bucket.Prefix().Scan(database),
166+
)
167+
})
168+
169+
t.Run("1 layer scan", func(t *testing.T) {
170+
for range testPerLayer {
171+
blockNumber, _ := randomEntry(t, content)
172+
validateResult(
173+
t,
174+
extractExpectedEntries(content, &blockNumber, nil, nil),
175+
bucket.Prefix().Add(blockNumber).Scan(database),
176+
)
177+
}
178+
})
179+
180+
t.Run("2 layer scan", func(t *testing.T) {
181+
for range testPerLayer * testPerLayer {
182+
blockNumber, map1 := randomEntry(t, content)
183+
hash, _ := randomEntry(t, map1)
184+
validateResult(
185+
t,
186+
extractExpectedEntries(content, &blockNumber, &hash, nil),
187+
bucket.Prefix().Add(blockNumber).Add(&hash).Scan(database),
188+
)
189+
}
190+
})
191+
192+
t.Run("3 layer scan", func(t *testing.T) {
193+
for range testPerLayer * testPerLayer * testPerLayer {
194+
blockNumber, map1 := randomEntry(t, content)
195+
hash, map2 := randomEntry(t, map1)
196+
slot, _ := randomEntry(t, map2)
197+
validateResult(
198+
t,
199+
extractExpectedEntries(content, &blockNumber, &hash, &slot),
200+
bucket.Prefix().Add(blockNumber).Add(&hash).Add([]byte(slot)).Scan(database),
201+
)
202+
}
203+
})
204+
}

‎db/typed/prefix/proxy.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package prefix
2+
3+
import "github.com/NethermindEth/juno/db/typed/key"
4+
5+
type proxy[V any, T tag[V]] struct{}
6+
7+
func Prefix[V any, H any, HS key.Serializer[H], T tag[V]](
8+
head HS,
9+
tail proxy[V, T],
10+
) proxy[V, hasPrefix[V, H, HS, T]] {
11+
return proxy[V, hasPrefix[V, H, HS, T]]{}
12+
}
13+
14+
func End[V any]() proxy[V, end[V]] {
15+
return proxy[V, end[V]]{}
16+
}

‎db/typed/prefix/scan_state.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package prefix
2+
3+
import (
4+
"iter"
5+
6+
"github.com/NethermindEth/juno/db"
7+
"github.com/NethermindEth/juno/db/dbutils"
8+
)
9+
10+
type scanState[V any] struct {
11+
prefix []byte
12+
scanner scanner[V]
13+
}
14+
15+
func (s scanState[V]) Scan(database db.KeyValueReader) iter.Seq2[Entry[V], error] {
16+
return s.scanner.scan(database, s.prefix)
17+
}
18+
19+
func (s scanState[V]) DeletePrefix(database db.KeyValueRangeDeleter) error {
20+
return database.DeleteRange(s.prefix, dbutils.UpperBound(s.prefix))
21+
}

‎db/typed/prefix/tag.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package prefix
2+
3+
import (
4+
"slices"
5+
6+
"github.com/NethermindEth/juno/db"
7+
"github.com/NethermindEth/juno/db/typed/key"
8+
)
9+
10+
type tag[V any] interface {
11+
~struct {
12+
scanState[V]
13+
}
14+
}
15+
16+
type hasPrefix[V any, H any, HS key.Serializer[H], T tag[V]] struct {
17+
scanState[V]
18+
}
19+
20+
func (h hasPrefix[V, H, HS, T]) Add(head H) T {
21+
return T{
22+
scanState: scanState[V]{
23+
prefix: append(h.prefix, HS{}.Marshal(head)...),
24+
scanner: h.scanner,
25+
},
26+
}
27+
}
28+
29+
func (h hasPrefix[V, H, HS, T]) DeleteRange(
30+
database db.KeyValueRangeDeleter,
31+
startKey,
32+
endKey H,
33+
) error {
34+
return database.DeleteRange(
35+
append(h.prefix, HS{}.Marshal(startKey)...),
36+
append(slices.Clone(h.prefix), HS{}.Marshal(endKey)...),
37+
)
38+
}
39+
40+
type end[V any] struct {
41+
scanState[V]
42+
}

0 commit comments

Comments
 (0)