Skip to content

Commit 37dd38d

Browse files
haikoscholhaikoschol
committed
use local KeysIter/PairsIter in StorageProvider
1 parent 2cd321d commit 37dd38d

File tree

3 files changed

+232
-27
lines changed

3 files changed

+232
-27
lines changed

‎internal/client/adapter/client_adapter.go

Lines changed: 46 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212

1313
"github.com/ChainSafe/gossamer/dot/state"
1414
"github.com/ChainSafe/gossamer/dot/types"
15+
"github.com/ChainSafe/gossamer/internal/client/api"
1516
"github.com/ChainSafe/gossamer/internal/database"
1617
"github.com/ChainSafe/gossamer/internal/primitives/blockchain"
1718
"github.com/ChainSafe/gossamer/internal/primitives/runtime"
@@ -789,34 +790,50 @@ func (ca *ClientAdapter[H, Hasher, N, E, Header]) StorageKeys(
789790
hash H,
790791
prefix,
791792
startKey storage.StorageKey,
792-
) (statemachine.KeysIter[H, Hasher], error) {
793+
) (api.KeysIter[H, Hasher], error) {
793794
stateAt, err := ca.client.StateAt(hash)
794795
if err != nil {
795-
return statemachine.KeysIter[H, Hasher]{}, err
796+
return api.KeysIter[H, Hasher]{}, err
797+
}
798+
backend, ok := stateAt.(*statemachine.TrieBackend[H, Hasher])
799+
if !ok {
800+
return api.KeysIter[H, Hasher]{}, fmt.Errorf(
801+
"got unexpected Backend type from StateAt() %T instead of statemachine.TrieBackend",
802+
backend,
803+
)
804+
}
805+
806+
iter, err := api.NewKeysIter(backend, &prefix, &startKey)
807+
if err != nil {
808+
return api.KeysIter[H, Hasher]{}, err
796809
}
797810

798-
return stateAt.Keys(statemachine.IterArgs{
799-
Prefix: prefix,
800-
StartAt: startKey,
801-
StartAtExclusive: true,
802-
})
811+
return *iter, nil
803812
}
804813

805814
func (ca *ClientAdapter[H, Hasher, N, E, Header]) StoragePairs(
806815
hash H,
807816
prefix,
808817
startKey storage.StorageKey,
809-
) (statemachine.PairsIter[H, Hasher], error) {
818+
) (api.PairsIter[H, Hasher], error) {
810819
stateAt, err := ca.client.StateAt(hash)
811820
if err != nil {
812-
return statemachine.PairsIter[H, Hasher]{}, err
821+
return api.PairsIter[H, Hasher]{}, err
822+
}
823+
backend, ok := stateAt.(*statemachine.TrieBackend[H, Hasher])
824+
if !ok {
825+
return api.PairsIter[H, Hasher]{}, fmt.Errorf(
826+
"got unexpected Backend type from StateAt() %T instead of statemachine.TrieBackend",
827+
backend,
828+
)
829+
}
830+
831+
iter, err := api.NewPairsIter(backend, &prefix, &startKey)
832+
if err != nil {
833+
return api.PairsIter[H, Hasher]{}, err
813834
}
814835

815-
return stateAt.Pairs(statemachine.IterArgs{
816-
Prefix: prefix,
817-
StartAt: startKey,
818-
StartAtExclusive: true,
819-
})
836+
return *iter, nil
820837
}
821838

822839
func (ca *ClientAdapter[H, Hasher, N, E, Header]) ChildStorage(
@@ -838,18 +855,25 @@ func (ca *ClientAdapter[H, Hasher, N, E, Header]) ChildStorageKeys(
838855
childInfo storage.ChildInfo,
839856
prefix storage.StorageKey,
840857
startKey storage.StorageKey,
841-
) (statemachine.KeysIter[H, Hasher], error) {
858+
) (api.KeysIter[H, Hasher], error) {
842859
stateAt, err := ca.client.StateAt(hash)
843860
if err != nil {
844-
return statemachine.KeysIter[H, Hasher]{}, err
861+
return api.KeysIter[H, Hasher]{}, err
862+
}
863+
backend, ok := stateAt.(*statemachine.TrieBackend[H, Hasher])
864+
if !ok {
865+
return api.KeysIter[H, Hasher]{}, fmt.Errorf(
866+
"got unexpected Backend type from StateAt() %T instead of statemachine.TrieBackend",
867+
backend,
868+
)
869+
}
870+
871+
iter, err := api.NewChildKeysIter(backend, childInfo, &prefix, &startKey)
872+
if err != nil {
873+
return api.KeysIter[H, Hasher]{}, err
845874
}
846875

847-
return stateAt.Keys(statemachine.IterArgs{
848-
Prefix: prefix,
849-
StartAt: startKey,
850-
StartAtExclusive: true,
851-
ChildInfo: childInfo,
852-
})
876+
return *iter, nil
853877
}
854878

855879
func (ca *ClientAdapter[H, Hasher, N, E, Header]) ChildStorageHash(

‎internal/client/adapter/client_adapter_test.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/ChainSafe/gossamer/dot/state"
1212
"github.com/ChainSafe/gossamer/dot/types"
1313
"github.com/ChainSafe/gossamer/internal/client/adapter/mocks"
14+
"github.com/ChainSafe/gossamer/internal/client/api"
1415
"github.com/ChainSafe/gossamer/internal/database"
1516
"github.com/ChainSafe/gossamer/internal/primitives/blockchain"
1617
"github.com/ChainSafe/gossamer/internal/primitives/core/hash"
@@ -59,6 +60,10 @@ var blockchainInfo = blockchain.Info[Hash, Number]{
5960
FinalizedNumber: blockNumber,
6061
}
6162

63+
func TestStorageProviderImplemented(t *testing.T) {
64+
var _ api.StorageProvider[Hash, Hasher] = &ClientAdapter[Hash, Hasher, Number, Extrinsic, Header]{}
65+
}
66+
6267
func TestBlockStateImplemented(t *testing.T) {
6368
var _ state.BlockState = &ClientAdapter[Hash, Hasher, Number, Extrinsic, Header]{}
6469
}

‎internal/client/api/backend.go

Lines changed: 181 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package api
55

66
import (
7+
"iter"
78
"sync"
89

910
"github.com/ChainSafe/gossamer/internal/primitives/blockchain"
@@ -352,6 +353,181 @@ type Backend[
352353
// UsageInfo() *UsageInfo
353354
}
354355

356+
// KeysIter is an iterator over storage keys.
357+
type KeysIter[H runtime.Hash, Hasher runtime.Hasher[H]] struct {
358+
backend *statemachine.TrieBackend[H, Hasher]
359+
rawIter statemachine.StorageIterator[H, Hasher]
360+
}
361+
362+
func NewKeysIter[H runtime.Hash, Hasher runtime.Hasher[H]](
363+
backend *statemachine.TrieBackend[H, Hasher],
364+
prefix *storage.StorageKey,
365+
startAt *storage.StorageKey,
366+
) (*KeysIter[H, Hasher], error) {
367+
args := statemachine.IterArgs{
368+
StartAtExclusive: true,
369+
}
370+
371+
if prefix != nil {
372+
args.Prefix = *prefix
373+
}
374+
375+
if startAt != nil {
376+
args.StartAt = *startAt
377+
}
378+
379+
rawIter, err := backend.RawIter(args)
380+
if err != nil {
381+
return nil, err
382+
}
383+
384+
return &KeysIter[H, Hasher]{backend, rawIter}, nil
385+
}
386+
387+
func NewChildKeysIter[H runtime.Hash, Hasher runtime.Hasher[H]](
388+
backend *statemachine.TrieBackend[H, Hasher],
389+
childInfo storage.ChildInfo,
390+
prefix *storage.StorageKey,
391+
startAt *storage.StorageKey,
392+
) (*KeysIter[H, Hasher], error) {
393+
args := statemachine.IterArgs{
394+
ChildInfo: childInfo,
395+
StartAtExclusive: true,
396+
}
397+
398+
if prefix != nil {
399+
args.Prefix = *prefix
400+
}
401+
402+
if startAt != nil {
403+
args.StartAt = *startAt
404+
}
405+
406+
rawIter, err := backend.RawIter(args)
407+
if err != nil {
408+
return nil, err
409+
}
410+
411+
return &KeysIter[H, Hasher]{backend, rawIter}, nil
412+
}
413+
414+
func (ki *KeysIter[H, Hasher]) Next() (storage.StorageKey, error) {
415+
key, err := ki.rawIter.NextKey(ki.backend)
416+
return storage.StorageKey(key), err
417+
}
418+
419+
func (ki *KeysIter[H, Hasher]) All() iter.Seq2[storage.StorageKey, error] {
420+
return func(yield func(storage.StorageKey, error) bool) {
421+
for {
422+
item, err := ki.Next()
423+
if err != nil {
424+
return
425+
}
426+
if item == nil {
427+
return
428+
}
429+
if !yield(item, err) {
430+
return
431+
}
432+
}
433+
}
434+
}
435+
436+
// StorageKeyData is a storage key/value pair.
437+
type StorageKeyData struct {
438+
storage.StorageKey
439+
storage.StorageData
440+
}
441+
442+
// PairsIter is an iterator over storage key/value pairs.
443+
type PairsIter[H runtime.Hash, Hasher runtime.Hasher[H]] struct {
444+
backend *statemachine.TrieBackend[H, Hasher]
445+
rawIter statemachine.StorageIterator[H, Hasher]
446+
}
447+
448+
func NewPairsIter[H runtime.Hash, Hasher runtime.Hasher[H]](
449+
backend *statemachine.TrieBackend[H, Hasher],
450+
prefix *storage.StorageKey,
451+
startAt *storage.StorageKey,
452+
) (*PairsIter[H, Hasher], error) {
453+
args := statemachine.IterArgs{
454+
StartAtExclusive: true,
455+
}
456+
457+
if prefix != nil {
458+
args.Prefix = *prefix
459+
}
460+
461+
if startAt != nil {
462+
args.StartAt = *startAt
463+
}
464+
465+
rawIter, err := backend.RawIter(args)
466+
if err != nil {
467+
return nil, err
468+
}
469+
470+
return &PairsIter[H, Hasher]{backend, rawIter}, nil
471+
}
472+
473+
func NewChildPairsIter[H runtime.Hash, Hasher runtime.Hasher[H]](
474+
backend *statemachine.TrieBackend[H, Hasher],
475+
childInfo storage.ChildInfo,
476+
prefix *storage.StorageKey,
477+
startAt *storage.StorageKey,
478+
) (*PairsIter[H, Hasher], error) {
479+
args := statemachine.IterArgs{
480+
ChildInfo: childInfo,
481+
StartAtExclusive: true,
482+
}
483+
484+
if prefix != nil {
485+
args.Prefix = *prefix
486+
}
487+
488+
if startAt != nil {
489+
args.StartAt = *startAt
490+
}
491+
492+
rawIter, err := backend.RawIter(args)
493+
if err != nil {
494+
return nil, err
495+
}
496+
497+
return &PairsIter[H, Hasher]{backend, rawIter}, nil
498+
}
499+
500+
func (ki *PairsIter[H, Hasher]) Next() (*StorageKeyData, error) {
501+
keyValue, err := ki.rawIter.NextKeyValue(ki.backend)
502+
if err != nil {
503+
return nil, err
504+
}
505+
506+
data := StorageKeyData{
507+
StorageKey: storage.StorageKey(keyValue.StorageKey),
508+
StorageData: storage.StorageData(keyValue.StorageValue),
509+
}
510+
511+
return &data, nil
512+
}
513+
514+
func (ki *PairsIter[H, Hasher]) All() iter.Seq2[StorageKeyData, error] {
515+
return func(yield func(StorageKeyData, error) bool) {
516+
for {
517+
item, err := ki.Next()
518+
if err != nil {
519+
return
520+
}
521+
if item == nil {
522+
return
523+
}
524+
if !yield(*item, err) {
525+
return
526+
}
527+
}
528+
}
529+
}
530+
355531
// StorageProvider provides access to storage primitives
356532
type StorageProvider[H runtime.Hash, Hasher runtime.Hasher[H]] interface {
357533
// Storage returns the value under the key in that block, given a blocks hash and a key.
@@ -360,13 +536,13 @@ type StorageProvider[H runtime.Hash, Hasher runtime.Hasher[H]] interface {
360536
// StorageHash returns the value under the hash in that block, given a blocks hash and a key.
361537
StorageHash(hash H, key storage.StorageKey) (*H, error)
362538

363-
// StorageKeys returns a [statemachine.KeysIter] that iterates over matching storage keys in that block
539+
// StorageKeys returns a [KeysIter] that iterates over matching storage keys in that block
364540
// given a blocks hash and a key prefix.
365-
StorageKeys(hash H, prefix, startKey storage.StorageKey) (statemachine.KeysIter[H, Hasher], error)
541+
StorageKeys(hash H, prefix, startKey storage.StorageKey) (KeysIter[H, Hasher], error)
366542

367543
// StoragePairs returns an iterator over the storage keys and values in that block,
368544
// given the blocks hash and a key prefix.
369-
StoragePairs(hash H, prefix, startKey storage.StorageKey) (statemachine.PairsIter[H, Hasher], error)
545+
StoragePairs(hash H, prefix, startKey storage.StorageKey) (PairsIter[H, Hasher], error)
370546

371547
// ChildStorage returns the value under the key in that block, given a blocks hash,
372548
// a key and a child storage key.
@@ -376,14 +552,14 @@ type StorageProvider[H runtime.Hash, Hasher runtime.Hasher[H]] interface {
376552
key storage.StorageKey,
377553
) (storage.StorageData, error)
378554

379-
// ChildStorageKeys returns a [statemachine.KeysIter] that iterates matching storage keys in that block,
555+
// ChildStorageKeys returns a [KeysIter] that iterates matching storage keys in that block,
380556
// given a blocks hash, an optional key prefix and an optional child storage key.
381557
ChildStorageKeys(
382558
hash H,
383559
childInfo storage.ChildInfo,
384560
prefix storage.StorageKey,
385561
startKey storage.StorageKey,
386-
) (statemachine.KeysIter[H, Hasher], error)
562+
) (KeysIter[H, Hasher], error)
387563

388564
// ChildStorageHash returns the hash under the key in a block, given its hash,
389565
// a key and a child storage key.

0 commit comments

Comments
 (0)