chore(trie): Add support for trie V1

cmd/gossamer/commands/import_state.go

@@ -7,13 +7,18 @@ import (
 	"fmt"
 
 	"github.com/ChainSafe/gossamer/dot"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/lib/utils"
 	"github.com/spf13/cobra"
 )
 
 func init() {
 	ImportStateCmd.Flags().String("chain", "", "Chain id used to load default configuration for specified chain")
 	ImportStateCmd.Flags().String("state-file", "", "Path to JSON file consisting of key-value pairs")
+	ImportStateCmd.Flags().Uint32("state-version",
+		uint32(trie.DefaultStateVersion),
+		"State version to use when importing state",
+	)
 	ImportStateCmd.Flags().String("header-file", "", "Path to JSON file of block header corresponding to the given state")
 	ImportStateCmd.Flags().Uint64("first-slot", 0, "The first BABE slot of the network")
 }
@@ -26,7 +31,8 @@ var ImportStateCmd = &cobra.Command{
 in the form of key-value pairs to be imported.
 Input can be generated by using the RPC function state_getPairs.
 Example: 
-	gossamer import-state --state-file state.json --header-file header.json --first-slot <first slot of network>`,
+	gossamer import-state --state-file state.json --state-version 1 --header-file header.json 
+	--first-slot <first slot of network>`,
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return execImportState(cmd)
 	},
@@ -54,6 +60,15 @@ func execImportState(cmd *cobra.Command) error {
 		return fmt.Errorf("state-file must be specified")
 	}
 
+	stateVersion, err := cmd.Flags().GetUint32("state-version")
+	if err != nil {
+		return fmt.Errorf("failed to get state-version: %s", err)
+	}
+	stateTrieVersion, err := trie.ParseVersion(stateVersion)
+	if err != nil {
+		return fmt.Errorf("invalid state version")
+	}
+
 	headerFile, err := cmd.Flags().GetString("header-file")
 	if err != nil {
 		return fmt.Errorf("failed to get header-file: %s", err)
@@ -64,5 +79,5 @@ func execImportState(cmd *cobra.Command) error {
 
 	basePath = utils.ExpandDir(basePath)
 
-	return dot.ImportState(basePath, stateFile, headerFile, firstSlot)
+	return dot.ImportState(basePath, stateFile, headerFile, stateTrieVersion, firstSlot)
 }

cmd/gossamer/commands/import_state_test.go

@@ -0,0 +1,57 @@
+// Copyright 2023 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package commands
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestImportStateMissingStateFile(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name()})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "state-file must be specified")
+}
+
+func TestImportStateInvalidFirstSlot(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(), "--first-slot", "wrong"})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "invalid argument \"wrong\"")
+}
+
+func TestImportStateEmptyHeaderFile(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(),
+		"--state-file", "test",
+		"--header-file", "",
+	})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "header-file must be specified")
+}
+
+func TestImportStateErrorImportingState(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(),
+		"--state-file", "test",
+		"--header-file", "test",
+	})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "no such file or directory")
+}

dot/core/helpers_test.go

@@ -57,7 +57,7 @@ func createTestService(t *testing.T, genesisFilePath string,
 	require.NoError(t, err)
 
 	genesisHeader := &types.Header{
-		StateRoot: genesisTrie.MustHash(),
+		StateRoot: trie.V0.MustHash(genesisTrie),
 		Number:    0,
 	}
 
@@ -271,7 +271,7 @@ func newWestendLocalWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash()
+	stateRoot := trie.V0.MustHash(genesisTrie)
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()

dot/core/service.go

@@ -20,6 +20,7 @@ import (
 	rtstorage "github.com/ChainSafe/gossamer/lib/runtime/storage"
 	wazero_runtime "github.com/ChainSafe/gossamer/lib/runtime/wazero"
 	"github.com/ChainSafe/gossamer/lib/transaction"
+	"github.com/ChainSafe/gossamer/lib/trie"
 
 	cscale "github.com/centrifuge/go-substrate-rpc-client/v4/scale"
 	ctypes "github.com/centrifuge/go-substrate-rpc-client/v4/types"
@@ -116,14 +117,34 @@ func (s *Service) Stop() error {
 	return nil
 }
 
+func (s *Service) getCurrentStateTrieVersion() (trie.TrieLayout, error) {
+	bestBlockHash := s.blockState.BestBlockHash()
+	rt, err := s.blockState.GetRuntime(bestBlockHash)
+	if err != nil {
+		return trie.NoVersion, err
+	}
+
+	runtimeVersion, err := rt.Version()
+	if err != nil {
+		return trie.NoVersion, err
+	}
+
+	return trie.ParseVersion(runtimeVersion.StateVersion)
+}
+
 // StorageRoot returns the hash of the storage root
 func (s *Service) StorageRoot() (common.Hash, error) {
 	ts, err := s.storageState.TrieState(nil)
 	if err != nil {
 		return common.Hash{}, err
 	}
 
-	return ts.Root()
+	stateTrieVersion, err := s.getCurrentStateTrieVersion()
+	if err != nil {
+		return common.Hash{}, err
+	}
+
+	return stateTrieVersion.Hash(ts.Trie())
 }
 
 // HandleBlockImport handles a block that was imported via the network
@@ -226,7 +247,7 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	}
 
 	logger.Debugf("imported block %s and stored state trie with root %s",
-		block.Header.Hash(), state.MustRoot())
+		block.Header.Hash(), state.MustRoot(trie.NoMaxInlineValueSize))
 
 	parentRuntimeInstance, err := s.blockState.GetRuntime(block.Header.ParentHash)
 	if err != nil {

dot/core/service_test.go

@@ -137,6 +137,7 @@ func Test_Service_StorageRoot(t *testing.T) {
 		retErr        error
 		expErr        error
 		expErrMsg     string
+		stateVersion  uint32
 	}{
 		{
 			name:          "storage trie state error",
@@ -147,12 +148,22 @@ func Test_Service_StorageRoot(t *testing.T) {
 			trieStateCall: true,
 		},
 		{
-			name:    "storage trie state ok",
+			name:    "storage trie state ok v0",
 			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:  ts,
 			trieStateCall: true,
+			stateVersion:  0,
+		},
+		{
+			name:    "storage trie state ok v1",
+			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:  ts,
+			trieStateCall: true,
+			stateVersion:  1,
 		},
 	}
 	for _, tt := range tests {
@@ -164,7 +175,23 @@ func Test_Service_StorageRoot(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				mockStorageState := NewMockStorageState(ctrl)
 				mockStorageState.EXPECT().TrieState(nil).Return(tt.retTrieState, tt.retErr)
+
 				service.storageState = mockStorageState
+
+				if tt.retErr == nil {
+					mockRuntimeVersion := runtime.Version{
+						StateVersion: tt.stateVersion,
+					}
+
+					mockRuntime := NewMockInstance(ctrl)
+					mockRuntime.EXPECT().Version().Return(mockRuntimeVersion, nil)
+
+					mockBlockState := NewMockBlockState(ctrl)
+					mockBlockState.EXPECT().BestBlockHash().Return(common.Hash{})
+					mockBlockState.EXPECT().GetRuntime(gomock.Any()).Return(mockRuntime, nil)
+
+					service.blockState = mockBlockState
+				}
 			}
 
 			res, err := service.StorageRoot()

dot/digest/helpers_test.go

@@ -28,7 +28,10 @@ func newWestendDevGenesisWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash()
+
+	// We are using state trie V0 since we are using the genesis trie where v0 is used
+	stateRoot := trie.V0.MustHash(genesisTrie)
+
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()

dot/helpers_test.go

@@ -42,7 +42,10 @@ func newWestendDevGenesisWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash()
+
+	// We are using state trie V0 since we are using the genesis trie where v0 is used
+	stateRoot := trie.V0.MustHash(genesisTrie)
+
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()

dot/import.go

@@ -20,7 +20,7 @@ import (
 )
 
 // ImportState imports the state in the given files to the database with the given path.
-func ImportState(basepath, stateFP, headerFP string, firstSlot uint64) error {
+func ImportState(basepath, stateFP, headerFP string, stateTrieVersion trie.TrieLayout, firstSlot uint64) error {
 	tr, err := newTrieFromPairs(stateFP)
 	if err != nil {
 		return err
@@ -38,7 +38,7 @@ func ImportState(basepath, stateFP, headerFP string, firstSlot uint64) error {
 		LogLevel: log.Info,
 	}
 	srv := state.NewService(config)
-	return srv.Import(header, tr, firstSlot)
+	return srv.Import(header, tr, stateTrieVersion, firstSlot)
 }
 
 func newTrieFromPairs(filename string) (*trie.Trie, error) {

dot/import_integration_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -22,20 +23,28 @@ func Test_newTrieFromPairs(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name     string
-		filename string
-		want     common.Hash
-		err      error
+		name         string
+		filename     string
+		want         common.Hash
+		stateVersion trie.TrieLayout
+		err          error
 	}{
 		{
 			name: "no_arguments",
 			err:  errors.New("read .: is a directory"),
 			want: common.Hash{},
 		},
 		{
-			name:     "working example",
-			filename: setupStateFile(t),
-			want:     common.MustHexToHash("0x09f9ca28df0560c2291aa16b56e15e07d1e1927088f51356d522722aa90ca7cb"),
+			name:         "working example",
+			filename:     setupStateFile(t),
+			want:         common.MustHexToHash("0x09f9ca28df0560c2291aa16b56e15e07d1e1927088f51356d522722aa90ca7cb"),
+			stateVersion: trie.V0,
+		},
+		{
+			name:         "working example",
+			filename:     setupStateFile(t),
+			want:         common.MustHexToHash("0xcc25fe024a58297658e576e2e4c33691fe3a9fe5a7cdd2e55534164a0fcc0782"),
+			stateVersion: trie.V1,
 		},
 	}
 	for _, tt := range tests {
@@ -52,7 +61,7 @@ func Test_newTrieFromPairs(t *testing.T) {
 			if tt.want.IsEmpty() {
 				assert.Nil(t, got)
 			} else {
-				assert.Equal(t, tt.want, got.MustHash())
+				assert.Equal(t, tt.want, tt.stateVersion.MustHash(*got))
 			}
 		})
 	}
@@ -93,7 +102,7 @@ func TestImportState_Integration(t *testing.T) {
 	headerFP := setupHeaderFile(t)
 
 	const firstSlot = uint64(262493679)
-	err = ImportState(config.BasePath, stateFP, headerFP, firstSlot)
+	err = ImportState(config.BasePath, stateFP, headerFP, trie.V0, firstSlot)
 	require.NoError(t, err)
 	// confirm data is imported into db
 	stateConfig := state.Config{
@@ -124,10 +133,11 @@ func TestImportState(t *testing.T) {
 	headerFP := setupHeaderFile(t)
 
 	type args struct {
-		basepath  string
-		stateFP   string
-		headerFP  string
-		firstSlot uint64
+		basepath     string
+		stateFP      string
+		headerFP     string
+		stateVersion trie.TrieLayout
+		firstSlot    uint64
 	}
 	tests := []struct {
 		name string
@@ -141,10 +151,11 @@ func TestImportState(t *testing.T) {
 		{
 			name: "working_example",
 			args: args{
-				basepath:  config.BasePath,
-				stateFP:   stateFP,
-				headerFP:  headerFP,
-				firstSlot: 262493679,
+				basepath:     config.BasePath,
+				stateFP:      stateFP,
+				headerFP:     headerFP,
+				stateVersion: trie.V0,
+				firstSlot:    262493679,
 			},
 		},
 	}
@@ -153,7 +164,7 @@ func TestImportState(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			err := ImportState(tt.args.basepath, tt.args.stateFP, tt.args.headerFP, tt.args.firstSlot)
+			err := ImportState(tt.args.basepath, tt.args.stateFP, tt.args.headerFP, tt.args.stateVersion, tt.args.firstSlot)
 			if tt.err != nil {
 				assert.EqualError(t, err, tt.err.Error())
 			} else {

dot/node_integration_test.go

@@ -388,7 +388,7 @@ func TestInitNode_LoadStorageRoot(t *testing.T) {
 	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"])
 	require.NoError(t, err)
 
-	expectedRoot, err := expected.Hash()
+	expectedRoot, err := trie.V0.Hash(&expected) // Since we are using a runtime with state trie V0
 	require.NoError(t, err)
 
 	coreServiceInterface := node.ServiceRegistry.Get(&core.Service{})

dot/rpc/helpers_test.go

@@ -28,7 +28,7 @@ func newWestendDevGenesisWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash()
+	stateRoot := trie.V0.MustHash(genesisTrie)
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()