refactor(pkg/trie): Encapsulate in-memory trie impl in its own package (

#3803) Co-authored-by: JimboJ <[email protected]>
ChainSafe · Mar 28, 2024 · 476115a · 476115a
1 parent 523cb5b
commit 476115a
Show file tree

Hide file tree

Showing 62 changed files with 916 additions and 875 deletions.
diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
@@ -82,9 +82,9 @@ jobs:
 
       - name: Trie memory test
         run: |
-          sed -i 's/const skip = true/const skip = false/g' ./pkg/trie/mem_test.go
+          sed -i 's/const skip = true/const skip = false/g' ./pkg/trie/inmemory/mem_test.go
           go test -run ^Test_Trie_MemoryUsage$ ./pkg/trie
-          sed -i 's/const skip = false/const skip = true/g' ./pkg/trie/mem_test.go
+          sed -i 's/const skip = false/const skip = true/g' ./pkg/trie/inmemory/mem_test.go
 
       - name: Test - Race
         run: make test-using-race-detector

diff --git a/dot/core/service_integration_test.go b/dot/core/service_integration_test.go
@@ -26,7 +26,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/transaction"
 	"github.com/ChainSafe/gossamer/lib/utils"
 	"github.com/ChainSafe/gossamer/pkg/scale"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 )
@@ -638,7 +638,7 @@ func TestService_HandleCodeSubstitutes(t *testing.T) {
 
 	s.blockState.StoreRuntime(blockHash, rt)
 
-	ts := rtstorage.NewTrieState(trie.NewEmptyTrie())
+	ts := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 	err = s.handleCodeSubstitution(blockHash, ts)
 	require.NoError(t, err)
 	codSub := s.codeSubstitutedState.(*state.BaseState).LoadCodeSubstitutedBlockHash()
@@ -666,7 +666,7 @@ func TestService_HandleRuntimeChangesAfterCodeSubstitutes(t *testing.T) {
 		Body: *body,
 	}
 
-	ts := rtstorage.NewTrieState(trie.NewEmptyTrie())
+	ts := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 	err = s.handleCodeSubstitution(blockHash, ts)
 	require.NoError(t, err)
 	require.Equal(t, codeHashBefore, parentRt.GetCodeHash()) // codeHash should remain unchanged after code substitute

diff --git a/dot/core/service_test.go b/dot/core/service_test.go
@@ -24,12 +24,13 @@ import (
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
 	"github.com/ChainSafe/gossamer/lib/transaction"
 	"github.com/ChainSafe/gossamer/pkg/scale"
-	"github.com/ChainSafe/gossamer/pkg/trie"
 	cscale "github.com/centrifuge/go-substrate-rpc-client/v4/scale"
 	"github.com/centrifuge/go-substrate-rpc-client/v4/signature"
 	ctypes "github.com/centrifuge/go-substrate-rpc-client/v4/types"
 	"github.com/centrifuge/go-substrate-rpc-client/v4/types/codec"
 
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
@@ -150,7 +151,7 @@ func Test_Service_StorageRoot(t *testing.T) {
 			service: &Service{},
 			exp: common.Hash{0x3, 0x17, 0xa, 0x2e, 0x75, 0x97, 0xb7, 0xb7, 0xe3, 0xd8, 0x4c, 0x5, 0x39, 0x1d, 0x13, 0x9a,
 				0x62, 0xb1, 0x57, 0xe7, 0x87, 0x86, 0xd8, 0xc0, 0x82, 0xf2, 0x9d, 0xcf, 0x4c, 0x11, 0x13, 0x14},
-			retTrieState:  rtstorage.NewTrieState(trie.NewEmptyTrie()),
+			retTrieState:  rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie()),
 			trieStateCall: true,
 			stateVersion:  0,
 		},
@@ -159,7 +160,7 @@ func Test_Service_StorageRoot(t *testing.T) {
 			service: &Service{},
 			exp: common.Hash{0x3, 0x17, 0xa, 0x2e, 0x75, 0x97, 0xb7, 0xb7, 0xe3, 0xd8, 0x4c, 0x5, 0x39, 0x1d, 0x13, 0x9a,
 				0x62, 0xb1, 0x57, 0xe7, 0x87, 0x86, 0xd8, 0xc0, 0x82, 0xf2, 0x9d, 0xcf, 0x4c, 0x11, 0x13, 0x14},
-			retTrieState:  rtstorage.NewTrieState(trie.NewEmptyTrie()),
+			retTrieState:  rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie()),
 			trieStateCall: true,
 			stateVersion:  1,
 		},
@@ -280,7 +281,7 @@ func Test_Service_handleCodeSubstitution(t *testing.T) {
 			errWrapped: io.ErrUnexpectedEOF,
 		},
 		"store_code_substitution_block_hash_error": {
-			trieState: rtstorage.NewTrieState(trie.NewEmptyTrie()),
+			trieState: rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie()),
 			serviceBuilder: func(ctrl *gomock.Controller) *Service {
 				storedRuntime := NewMockInstance(ctrl)
 				storedRuntime.EXPECT().Keystore().Return(nil)
@@ -338,7 +339,7 @@ func Test_Service_handleCodeSubstitution(t *testing.T) {
 				}
 			},
 			blockHash: common.Hash{0x01},
-			trieState: rtstorage.NewTrieState(trie.NewEmptyTrie()),
+			trieState: rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie()),
 		},
 	}
 
@@ -378,7 +379,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("storeTrie_error", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -394,7 +395,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("addBlock_quit_error", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -415,7 +416,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("addBlock_parent_not_found_error", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -436,7 +437,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("addBlock_error_continue", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -465,7 +466,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("handle_runtime_changes_error", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -497,7 +498,7 @@ func Test_Service_handleBlock(t *testing.T) {
 
 	t.Run("code_substitution_ok", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		testHeader := types.NewEmptyHeader()
 		block := types.NewBlock(*testHeader, *types.NewBody([]types.Extrinsic{[]byte{21}}))
@@ -544,7 +545,7 @@ func Test_Service_HandleBlockProduced(t *testing.T) {
 
 	t.Run("happy_path", func(t *testing.T) {
 		t.Parallel()
-		trieState := rtstorage.NewTrieState(trie.NewEmptyTrie())
+		trieState := rtstorage.NewTrieState(inmemory_trie.NewEmptyTrie())
 
 		digest := types.NewDigest()
 		err := digest.Add(
@@ -1144,7 +1145,7 @@ func TestServiceGetRuntimeVersion(t *testing.T) {
 		}},
 		TransactionVersion: transactionVersion,
 	}
-	emptyTrie := trie.NewEmptyTrie()
+	emptyTrie := inmemory_trie.NewEmptyTrie()
 	ts := rtstorage.NewTrieState(emptyTrie)
 
 	execTest := func(t *testing.T, s *Service, bhash *common.Hash, exp runtime.Version,

diff --git a/dot/import.go b/dot/import.go
@@ -15,6 +15,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 
 	"github.com/ChainSafe/gossamer/internal/log"
 )
@@ -62,7 +63,7 @@ func newTrieFromPairs(filename string, version trie.TrieLayout) (trie.Trie, erro
 		entries[pairArr[0].(string)] = pairArr[1].(string)
 	}
 
-	tr, err := trie.LoadFromMap(entries, version)
+	tr, err := inmemory_trie.LoadFromMap(entries, version)
 	if err != nil {
 		return nil, err
 	}

diff --git a/dot/node_integration_test.go b/dot/node_integration_test.go
@@ -36,6 +36,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
 	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	gomock "go.uber.org/mock/gomock"
@@ -385,7 +386,7 @@ func TestInitNode_LoadStorageRoot(t *testing.T) {
 	node, err := NewNode(config, ks)
 	require.NoError(t, err)
 
-	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"], trie.V0)
+	expected, err := inmemory_trie.LoadFromMap(gen.GenesisFields().Raw["top"], trie.V0)
 	require.NoError(t, err)
 
 	expectedRoot, err := trie.V0.Hash(expected) // Since we are using a runtime with state trie V0

diff --git a/dot/rpc/modules/childstate_integration_test.go b/dot/rpc/modules/childstate_integration_test.go
@@ -13,7 +13,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/database"
 	"github.com/ChainSafe/gossamer/lib/common"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/stretchr/testify/require"
 )
 
@@ -91,7 +91,8 @@ func TestChildStateGetStorageSize(t *testing.T) {
 			keyChild: []byte(":child_storage_key"),
 		},
 		{
-			err:      fmt.Errorf("child trie does not exist at key 0x%x%x", trie.ChildStorageKeyPrefix, []byte(":not_exist")),
+			err: fmt.Errorf("child trie does not exist at key 0x%x%x",
+				inmemory_trie.ChildStorageKeyPrefix, []byte(":not_exist")),
 			hash:     &blockHash,
 			entry:    []byte(":child_second"),
 			keyChild: []byte(":not_exist"),
@@ -149,7 +150,7 @@ func TestGetStorageHash(t *testing.T) {
 		},
 		{
 			err: fmt.Errorf("child trie does not exist at key 0x%x%x",
-				string(trie.ChildStorageKeyPrefix), []byte(":not_exist")),
+				string(inmemory_trie.ChildStorageKeyPrefix), []byte(":not_exist")),
 			hash:     &blockHash,
 			entry:    []byte(":child_second"),
 			keyChild: []byte(":not_exist"),

diff --git a/dot/rpc/modules/dev_integration_test.go b/dot/rpc/modules/dev_integration_test.go
@@ -17,7 +17,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/keystore"
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 )
@@ -56,7 +56,7 @@ func newBABEService(t *testing.T) *babe.Service {
 	require.NoError(t, err)
 
 	bs, es := newState(t)
-	tt := trie.NewEmptyTrie()
+	tt := inmemory_trie.NewEmptyTrie()
 	rt := wazero_runtime.NewTestInstance(t, runtime.WESTEND_RUNTIME_v0929, wazero_runtime.TestWithTrie(tt))
 	bs.StoreRuntime(bs.GenesisHash(), rt)
 	tt.Put(

diff --git a/dot/rpc/modules/system_integration_test.go b/dot/rpc/modules/system_integration_test.go
@@ -30,7 +30,7 @@ import (
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
 	"github.com/ChainSafe/gossamer/lib/transaction"
 	"github.com/ChainSafe/gossamer/pkg/scale"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 )
 
 var (
@@ -353,7 +353,7 @@ func setupSystemModule(t *testing.T) *SystemModule {
 
 func newCoreService(t *testing.T, srvc *state.Service) *core.Service {
 	// setup service
-	tt := trie.NewEmptyTrie()
+	tt := inmemory_trie.NewEmptyTrie()
 	rt := wazero_runtime.NewTestInstance(t, runtime.WESTEND_RUNTIME_v0929, wazero_runtime.TestWithTrie(tt))
 	ks := keystore.NewGlobalKeystore()
 	t.Cleanup(func() {

diff --git a/dot/state/base_test.go b/dot/state/base_test.go
@@ -8,14 +8,14 @@ import (
 
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/genesis"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestTrie_StoreAndLoadFromDB(t *testing.T) {
 	db := NewInMemoryDB(t)
-	tt := trie.NewEmptyTrie()
+	tt := inmemory_trie.NewEmptyTrie()
 
 	generator := newGenerator()
 	const size = 500
@@ -34,7 +34,7 @@ func TestTrie_StoreAndLoadFromDB(t *testing.T) {
 
 	expected := tt.MustHash()
 
-	tt = trie.NewEmptyTrie()
+	tt = inmemory_trie.NewEmptyTrie()
 	err = tt.Load(db, encroot)
 	require.NoError(t, err)
 	require.Equal(t, expected, tt.MustHash())

diff --git a/dot/state/block_finalisation_test.go b/dot/state/block_finalisation_test.go
@@ -8,7 +8,7 @@ import (
 
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 
 	"github.com/stretchr/testify/require"
 )
@@ -97,7 +97,7 @@ func TestBlockState_SetFinalisedHash(t *testing.T) {
 	require.NoError(t, err)
 
 	// set tries with some state root
-	bs.tries.softSet(someStateRoot, trie.NewEmptyTrie())
+	bs.tries.softSet(someStateRoot, inmemory_trie.NewEmptyTrie())
 
 	err = bs.SetFinalisedHash(testhash, 1, 1)
 	require.NoError(t, err)

diff --git a/dot/state/block_test.go b/dot/state/block_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"go.uber.org/mock/gomock"
 
 	"github.com/stretchr/testify/require"
@@ -40,7 +41,7 @@ func newTestBlockState(t *testing.T, tries *Tries) *BlockState {
 
 	// loads in-memory tries with genesis state root, should be deleted
 	// after another block is finalised
-	tr := trie.NewEmptyTrie()
+	tr := inmemory_trie.NewEmptyTrie()
 	err = tr.Load(bs.db, header.StateRoot)
 	require.NoError(t, err)
 	bs.tries.softSet(header.StateRoot, tr)

diff --git a/dot/state/grandpa_test.go b/dot/state/grandpa_test.go
@@ -15,7 +15,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/crypto/ed25519"
 	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
 	"github.com/ChainSafe/gossamer/lib/keystore"
-	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/gtank/merlin"
 	"go.uber.org/mock/gomock"
 
@@ -135,7 +135,7 @@ func testBlockState(t *testing.T, db database.Database) *BlockState {
 
 	// loads in-memory tries with genesis state root, should be deleted
 	// after another block is finalised
-	tr := trie.NewEmptyTrie()
+	tr := inmemory_trie.NewEmptyTrie()
 	err = tr.Load(bs.db, header.StateRoot)
 	require.NoError(t, err)
 	bs.tries.softSet(header.StateRoot, tr)

diff --git a/dot/state/initialize.go b/dot/state/initialize.go
@@ -15,6 +15,7 @@ import (
 	rtstorage "github.com/ChainSafe/gossamer/lib/runtime/storage"
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
 	"github.com/ChainSafe/gossamer/pkg/trie"
+	inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 )
 
 // Initialise initialises the genesis state of the DB using the given storage trie.
@@ -41,7 +42,7 @@ func (s *Service) Initialise(gen *genesis.Genesis, header *types.Header, t trie.
 	s.db = db
 
 	// TODO: all trie related db operations should be done in pkg/trie
-	if inmemoryTrie, ok := t.(*trie.InMemoryTrie); ok {
+	if inmemoryTrie, ok := t.(*inmemory_trie.InMemoryTrie); ok {
 		if err = inmemoryTrie.WriteDirty(database.NewTable(db, storagePrefix)); err != nil {
 			return fmt.Errorf("failed to write genesis trie to database: %w", err)
 		}
@@ -140,7 +141,7 @@ func loadGrandpaAuthorities(t trie.Trie) ([]types.GrandpaVoter, error) {
 func (s *Service) storeInitialValues(data *genesis.Data, t trie.Trie) error {
 	// write genesis trie to database
 	// TODO: all trie related db operations should be done in pkg/trie
-	if inmemoryTrie, ok := t.(*trie.InMemoryTrie); ok {
+	if inmemoryTrie, ok := t.(*inmemory_trie.InMemoryTrie); ok {
 		if err := inmemoryTrie.WriteDirty(database.NewTable(s.db, storagePrefix)); err != nil {
 			return fmt.Errorf("failed to write genesis trie to database: %w", err)
 		}