chore(trie): Add support for trie V1

cmd/gossamer/commands/import_state.go

@@ -7,13 +7,14 @@ 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().String("state-version", trie.DefaultStateVersion.String(), "State version v0 or v1")
 	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 +27,7 @@ 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 v1 --header-file header.json --first-slot <first slot of network>`, //nolint:lll
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return execImportState(cmd)
 	},
@@ -62,7 +63,16 @@ func execImportState(cmd *cobra.Command) error {
 		return fmt.Errorf("header-file must be specified")
 	}
 
+	stateVersionFlag, err := cmd.Flags().GetString("state-version")
+	if err != nil {
+		return fmt.Errorf("failed to get state-version: %s", err)
+	}
+	stateVersion, err := trie.ParseVersion(stateVersionFlag)
+	if err != nil {
+		return fmt.Errorf("failed to parse state-version: %s", err)
+	}
+
 	basePath = utils.ExpandDir(basePath)
 
-	return dot.ImportState(basePath, stateFile, headerFile, firstSlot)
+	return dot.ImportState(basePath, stateFile, headerFile, firstSlot, stateVersion)
 }

cmd/gossamer/commands/import_state_test.go

@@ -0,0 +1,83 @@
+// 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 TestImportStateEmptyStateFile(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(), "--state-version", "v0", "--state-file", ""})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "state-file must be specified")
+}
+
+func TestImportStateEmptyHeaderFile(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(),
+		"--state-version", "v0",
+		"--state-file", "test",
+		"--header-file", "",
+	})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "header-file must be specified")
+}
+
+func TestImportStateInvalidStateVersion(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(),
+		"--state-version", "v999",
+		"--state-file", "test",
+		"--header-file", "test",
+	})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "failed to parse state-version: parsing version failed: \"v999\" must be one of [v0, v1]")
+}
+
+func TestImportStateErrorImportingState(t *testing.T) {
+	rootCmd, err := NewRootCommand()
+	require.NoError(t, err)
+	rootCmd.AddCommand(ImportStateCmd)
+
+	rootCmd.SetArgs([]string{ImportStateCmd.Name(),
+		"--state-version", "v0",
+		"--state-file", "test",
+		"--header-file", "test",
+	})
+	err = rootCmd.Execute()
+	assert.ErrorContains(t, err, "no such file or directory")
+}

dot/core/helpers_test.go

@@ -109,9 +109,9 @@ func createTestService(t *testing.T, genesisFilePath string,
 	cfgRuntime, err := wazero_runtime.NewRuntimeFromGenesis(rtCfg)
 	require.NoError(t, err)
 
-	cfgRuntime.Context.Storage.Put(aliceBalanceKey, encodedAccountInfo)
+	cfgRuntime.Context.Storage.Put(aliceBalanceKey, encodedAccountInfo, trie.V0)
 	// this key is System.UpgradedToDualRefCount -> set to true since all accounts have been upgraded to v0.9 format
-	cfgRuntime.Context.Storage.Put(common.UpgradedToDualRefKey, []byte{1})
+	cfgRuntime.Context.Storage.Put(common.UpgradedToDualRefKey, []byte{1}, trie.V0)
 
 	cfgBlockState.StoreRuntime(cfgBlockState.BestBlockHash(), cfgRuntime)
 

dot/core/service_integration_test.go

@@ -584,7 +584,7 @@ func createBlockUsingNewRuntime(t *testing.T, bestBlockHash common.Hash, newRunt
 	testRuntime, err := os.ReadFile(newRuntimePath)
 	require.NoError(t, err)
 
-	trieState.Put(common.CodeKey, testRuntime)
+	trieState.Put(common.CodeKey, testRuntime, trie.V0)
 
 	primaryDigestData := types.NewBabePrimaryPreDigest(0, uint64(1), [32]byte{}, [64]byte{})
 	digest := types.NewDigest()
@@ -673,7 +673,7 @@ func TestService_HandleRuntimeChangesAfterCodeSubstitutes(t *testing.T) {
 	ts, err = s.storageState.TrieState(nil)
 	require.NoError(t, err)
 
-	ts.Put(common.CodeKey, testRuntime)
+	ts.Put(common.CodeKey, testRuntime, trie.V0)
 	rtUpdateBhash := newBlock.Header.Hash()
 
 	// update runtime for new block

dot/import.go

@@ -20,8 +20,8 @@ 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 {
-	tr, err := newTrieFromPairs(stateFP)
+func ImportState(basepath, stateFP, headerFP string, firstSlot uint64, stateVersion trie.Version) error {
+	tr, err := newTrieFromPairs(stateFP, stateVersion)
 	if err != nil {
 		return err
 	}
@@ -41,7 +41,7 @@ func ImportState(basepath, stateFP, headerFP string, firstSlot uint64) error {
 	return srv.Import(header, tr, firstSlot)
 }
 
-func newTrieFromPairs(filename string) (*trie.Trie, error) {
+func newTrieFromPairs(filename string, stateVersion trie.Version) (*trie.Trie, error) {
 	data, err := os.ReadFile(filepath.Clean(filename))
 	if err != nil {
 		return nil, err
@@ -62,7 +62,7 @@ func newTrieFromPairs(filename string) (*trie.Trie, error) {
 		entries[pairArr[0].(string)] = pairArr[1].(string)
 	}
 
-	tr, err := trie.LoadFromMap(entries)
+	tr, err := trie.LoadFromMap(entries, stateVersion)
 	if err != nil {
 		return nil, err
 	}

dot/import_integration_test.go

@@ -13,51 +13,12 @@ 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"
 )
 
-func Test_newTrieFromPairs(t *testing.T) {
-	t.Parallel()
-
-	tests := []struct {
-		name     string
-		filename string
-		want     common.Hash
-		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"),
-		},
-	}
-	for _, tt := range tests {
-		tt := tt
-		t.Run(tt.name, func(t *testing.T) {
-			t.Parallel()
-
-			got, err := newTrieFromPairs(tt.filename)
-			if tt.err != nil {
-				assert.EqualError(t, err, tt.err.Error())
-			} else {
-				assert.NoError(t, err)
-			}
-			if tt.want.IsEmpty() {
-				assert.Nil(t, got)
-			} else {
-				assert.Equal(t, tt.want, got.MustHash())
-			}
-		})
-	}
-}
-
 func TestNewHeaderFromFile(t *testing.T) {
 	fp := setupHeaderFile(t)
 	header, err := newHeaderFromFile(fp)
@@ -93,7 +54,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, firstSlot, trie.V0)
 	require.NoError(t, err)
 	// confirm data is imported into db
 	stateConfig := state.Config{
@@ -126,6 +87,7 @@ func TestImportState(t *testing.T) {
 	type args struct {
 		basepath  string
 		stateFP   string
+		version   trie.Version
 		headerFP  string
 		firstSlot uint64
 	}
@@ -153,7 +115,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.firstSlot, tt.args.version)
 			if tt.err != nil {
 				assert.EqualError(t, err, tt.err.Error())
 			} else {

dot/import_test.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/lib/utils"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -139,3 +140,45 @@ func Test_newHeaderFromFile(t *testing.T) {
 		})
 	}
 }
+
+func Test_newTrieFromPairs(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		filename string
+		version  trie.Version
+		want     common.Hash
+		err      error
+	}{
+		{
+			name: "no_arguments",
+			err:  errors.New("read .: is a directory"),
+			want: common.Hash{},
+		},
+		{
+			name:     "working example",
+			filename: setupStateFile(t),
+			version:  trie.V0,
+			want:     common.MustHexToHash("0x09f9ca28df0560c2291aa16b56e15e07d1e1927088f51356d522722aa90ca7cb"),
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := newTrieFromPairs(tt.filename, tt.version)
+			if tt.err != nil {
+				assert.EqualError(t, err, tt.err.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+			if tt.want.IsEmpty() {
+				assert.Nil(t, got)
+			} else {
+				assert.Equal(t, tt.want, got.MustHash())
+			}
+		})
+	}
+}

dot/node_integration_test.go

@@ -385,7 +385,7 @@ func TestInitNode_LoadStorageRoot(t *testing.T) {
 	node, err := NewNode(config, ks)
 	require.NoError(t, err)
 
-	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"])
+	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"], trie.V0)
 	require.NoError(t, err)
 
 	expectedRoot, err := expected.Hash()

dot/rpc/modules/author_integration_test.go

@@ -69,7 +69,7 @@ func useInstanceFromRuntimeV0929(t *testing.T, rtStorage *storage.TrieState) (in
 	bytes, err := os.ReadFile(testRuntimeFilePath)
 	require.NoError(t, err)
 
-	rtStorage.Put(common.CodeKey, bytes)
+	rtStorage.Put(common.CodeKey, bytes, trie.V0)
 
 	cfg := wazero_runtime.Config{
 		Role:     0,

dot/rpc/modules/childstate_integration_test.go

@@ -244,15 +244,15 @@ func setupChildStateStorage(t *testing.T) (*ChildStateModule, common.Hash) {
 	tr, err := st.Storage.TrieState(nil)
 	require.NoError(t, err)
 
-	tr.Put([]byte(":first_key"), []byte(":value1"))
-	tr.Put([]byte(":second_key"), []byte(":second_value"))
+	tr.Put([]byte(":first_key"), []byte(":value1"), trie.V0)
+	tr.Put([]byte(":second_key"), []byte(":second_value"), trie.V0)
 
 	childTr := trie.NewEmptyTrie()
-	childTr.Put([]byte(":child_first"), []byte(":child_first_value"))
-	childTr.Put([]byte(":child_second"), []byte(":child_second_value"))
-	childTr.Put([]byte(":another_child"), []byte("value"))
+	childTr.Put([]byte(":child_first"), []byte(":child_first_value"), trie.V0)
+	childTr.Put([]byte(":child_second"), []byte(":child_second_value"), trie.V0)
+	childTr.Put([]byte(":another_child"), []byte("value"), trie.V0)
 
-	err = tr.SetChild([]byte(":child_storage_key"), childTr)
+	err = tr.SetChild([]byte(":child_storage_key"), childTr, trie.V0)
 	require.NoError(t, err)
 
 	stateRoot, err := tr.Root()

dot/rpc/modules/childstate_test.go

@@ -24,15 +24,15 @@ func createTestTrieState(t *testing.T) (*trie.Trie, common.Hash) {
 	_, genesisTrie, _ := newWestendLocalGenesisWithTrieAndHeader(t)
 	tr := rtstorage.NewTrieState(&genesisTrie)
 
-	tr.Put([]byte(":first_key"), []byte(":value1"))
-	tr.Put([]byte(":second_key"), []byte(":second_value"))
+	tr.Put([]byte(":first_key"), []byte(":value1"), trie.V0)
+	tr.Put([]byte(":second_key"), []byte(":second_value"), trie.V0)
 
 	childTr := trie.NewEmptyTrie()
-	childTr.Put([]byte(":child_first"), []byte(":child_first_value"))
-	childTr.Put([]byte(":child_second"), []byte(":child_second_value"))
-	childTr.Put([]byte(":another_child"), []byte("value"))
+	childTr.Put([]byte(":child_first"), []byte(":child_first_value"), trie.V0)
+	childTr.Put([]byte(":child_second"), []byte(":child_second_value"), trie.V0)
+	childTr.Put([]byte(":another_child"), []byte("value"), trie.V0)
 
-	err := tr.SetChild([]byte(":child_storage_key"), childTr)
+	err := tr.SetChild([]byte(":child_storage_key"), childTr, trie.V0)
 	require.NoError(t, err)
 
 	stateRoot, err := tr.Root()

dot/rpc/modules/dev_integration_test.go

@@ -68,7 +68,8 @@ func newBABEService(t *testing.T) *babe.Service {
 			"0d4a9e054df4e01000000000000001cbd2d43530a44705ad088af313e18f80b53ef16b36177cd4b77b846f2a5f07c010000"+
 			"00000000004603307f855321776922daeea21ee31720388d097cdaac66f05a6f8462b317570100000000000000be1d9d59d"+
 			"e1283380100550a7b024501cb62d6cc40e3db35fcc5cf341814986e01000000000000001206960f920a23f7f4c43cc9081"+
-			"ec2ed0721f31a9bef2c10fd7602e16e08a32c0100000000000000"))
+			"ec2ed0721f31a9bef2c10fd7602e16e08a32c0100000000000000"),
+		trie.V0)
 
 	cfg := &babe.ServiceConfig{
 		BlockState:         bs,