Production code changes

dot/build_spec.go

@@ -112,8 +112,19 @@ func BuildFromDB(path string) (*BuildSpec, error) {
 	if err != nil {
 		return nil, fmt.Errorf("cannot start state service: %w", err)
 	}
+
+	bestBlockStateRoot, err := stateSrvc.Block.BestBlockStateRoot()
+	if err != nil {
+		return nil, fmt.Errorf("getting best block state root: %w", err)
+	}
+
+	stateVersion, err := stateSrvc.Block.GetRuntimeStateVersion(bestBlockStateRoot)
+	if err != nil {
+		return nil, fmt.Errorf("getting runtime state version for block state root: %w", err)
+	}
+
 	// set genesis fields data
-	ent, err := stateSrvc.Storage.Entries(nil)
+	ent, err := stateSrvc.Storage.Entries(&bestBlockStateRoot, stateVersion)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get storage trie entries: %w", err)
 	}

dot/core/interface.go

@@ -16,6 +16,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	rtstorage "github.com/ChainSafe/gossamer/lib/runtime/storage"
 	"github.com/ChainSafe/gossamer/lib/transaction"
+	"github.com/ChainSafe/gossamer/lib/trie"
 )
 
 // RuntimeInstance for runtime methods
@@ -30,6 +31,7 @@ type RuntimeInstance interface {
 	SetContextStorage(s runtime.Storage)
 	GetCodeHash() common.Hash
 	Version() (version runtime.Version)
+	StateVersion() (stateVersion trie.Version, err error)
 	Metadata() ([]byte, error)
 	BabeConfiguration() (*types.BabeConfiguration, error)
 	GrandpaAuthorities() ([]types.Authority, error)
@@ -70,19 +72,22 @@ type BlockState interface {
 	SubChain(start, end common.Hash) ([]common.Hash, error)
 	GetBlockBody(hash common.Hash) (*types.Body, error)
 	HandleRuntimeChanges(newState *rtstorage.TrieState, in runtime.Instance, bHash common.Hash) error
-	GetRuntime(*common.Hash) (runtime.Instance, error)
+	GetRuntime(blockHash *common.Hash) (instance runtime.Instance, err error)
+	GetRuntimeStateVersion(blockHash common.Hash) (stateVersion trie.Version, err error)
 	StoreRuntime(common.Hash, runtime.Instance)
 }
 
 // StorageState interface for storage state methods
 type StorageState interface {
 	LoadCode(root *common.Hash) ([]byte, error)
 	LoadCodeHash(root *common.Hash) (common.Hash, error)
-	TrieState(root *common.Hash) (*rtstorage.TrieState, error)
-	StoreTrie(*rtstorage.TrieState, *types.Header) error
+	TrieState(root *common.Hash, stateVersion trie.Version) (*rtstorage.TrieState, error)
+	StoreTrie(trieState *rtstorage.TrieState, blockHeader *types.Header,
+		stateVersion trie.Version) error
 	GetStateRootFromBlock(bhash *common.Hash) (*common.Hash, error)
 	GetStorage(root *common.Hash, key []byte) ([]byte, error)
-	GenerateTrieProof(stateRoot common.Hash, keys [][]byte) ([][]byte, error)
+	GenerateTrieProof(stateRoot common.Hash, keys [][]byte,
+		stateVersion trie.Version) (encodedNodes [][]byte, err error)
 	sync.Locker
 }
 

dot/core/messages.go

@@ -20,7 +20,14 @@ func (s *Service) validateTransaction(peerID peer.ID, head *types.Header, rt Run
 	tx types.Extrinsic) (validity *transaction.Validity, valid bool, err error) {
 	s.storageState.Lock()
 
-	ts, err := s.storageState.TrieState(&head.StateRoot)
+	// Note this is a cheap call getting the runtime cached version
+	// so we can call this in this function and not pass it as argument.
+	stateVersion, err := rt.StateVersion()
+	if err != nil {
+		return nil, false, fmt.Errorf("getting runtime state version: %w", err)
+	}
+
+	ts, err := s.storageState.TrieState(&head.StateRoot, stateVersion)
 	s.storageState.Unlock()
 	if err != nil {
 		return nil, false, fmt.Errorf("cannot get trie state from storage for root %s: %w", head.StateRoot, err)

dot/core/service.go

@@ -22,6 +22,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/runtime/wasmer"
 	"github.com/ChainSafe/gossamer/lib/services"
 	"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"
 )
@@ -60,6 +61,7 @@ type Service struct {
 
 // Config holds the configuration for the core Service.
 type Config struct {
+	// TODO add state version field here
 	LogLvl log.Level
 
 	BlockState       BlockState
@@ -102,7 +104,7 @@ func NewService(cfg *Config) (*Service, error) {
 	blockAddCh := make(chan *types.Block, 256)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	srv := &Service{
+	return &Service{
 		ctx:                  ctx,
 		cancel:               cancel,
 		keys:                 cfg.Keystore,
@@ -114,9 +116,7 @@ func NewService(cfg *Config) (*Service, error) {
 		blockAddCh:           blockAddCh,
 		codeSubstitute:       cfg.CodeSubstitutes,
 		codeSubstitutedState: cfg.CodeSubstitutedState,
-	}
-
-	return srv, nil
+	}, nil
 }
 
 // Start starts the core service
@@ -135,18 +135,21 @@ func (s *Service) Stop() error {
 	return nil
 }
 
-// StorageRoot returns the hash of the storage root
-func (s *Service) StorageRoot() (common.Hash, error) {
+// StorageRoot returns the hash of the storage root.
+// It is only used by tests, but has to exported because
+// internal fields are not exported to other packages.
+func (s *Service) StorageRoot(stateVersion trie.Version) (
+	rootHash common.Hash, err error) {
 	if s.storageState == nil {
-		return common.Hash{}, ErrNilStorageState
+		return rootHash, ErrNilStorageState
 	}
 
-	ts, err := s.storageState.TrieState(nil)
+	ts, err := s.storageState.TrieState(nil, stateVersion)
 	if err != nil {
-		return common.Hash{}, err
+		return rootHash, err
 	}
 
-	return ts.Root()
+	return ts.Root(stateVersion)
 }
 
 // HandleBlockImport handles a block that was imported via the network
@@ -188,8 +191,18 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 		return ErrNilBlockHandlerParameter
 	}
 
+	rt, err := s.blockState.GetRuntime(&block.Header.ParentHash)
+	if err != nil {
+		return fmt.Errorf("getting runtime: %w", err)
+	}
+
+	stateVersion, err := rt.StateVersion()
+	if err != nil {
+		return fmt.Errorf("getting runtime state version: %w", err)
+	}
+
 	// store updates state trie nodes in database
-	err := s.storageState.StoreTrie(state, &block.Header)
+	err = s.storageState.StoreTrie(state, &block.Header, stateVersion)
 	if err != nil {
 		logger.Warnf("failed to store state trie for imported block %s: %s",
 			block.Header.Hash(), err)
@@ -208,12 +221,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())
-
-	rt, err := s.blockState.GetRuntime(&block.Header.ParentHash)
-	if err != nil {
-		return err
-	}
+		block.Header.Hash(), state.MustRoot(stateVersion))
 
 	// check for runtime changes
 	if err := s.blockState.HandleRuntimeChanges(state, rt, block.Header.Hash()); err != nil {
@@ -467,14 +475,21 @@ func (s *Service) GetRuntimeVersion(bhash *common.Hash) (
 		}
 	}
 
-	ts, err := s.storageState.TrieState(stateRootHash)
+	rt, err := s.blockState.GetRuntime(bhash)
 	if err != nil {
 		return version, err
 	}
 
-	rt, err := s.blockState.GetRuntime(bhash)
+	stateVersion, err := rt.StateVersion()
 	if err != nil {
-		return version, err
+		return version, fmt.Errorf("getting runtime state version: %w", err)
+	}
+
+	// Note: not too sure why this trie state call is needed,
+	// but some RPC tests fail without it.
+	ts, err := s.storageState.TrieState(stateRootHash, stateVersion)
+	if err != nil {
+		return version, fmt.Errorf("getting trie state: %w", err)
 	}
 
 	rt.SetContextStorage(ts)
@@ -498,15 +513,19 @@ func (s *Service) HandleSubmittedExtrinsic(ext types.Extrinsic) error {
 		return fmt.Errorf("could not get state root from block %s: %w", bestBlockHash, err)
 	}
 
-	ts, err := s.storageState.TrieState(stateRoot)
+	rt, err := s.blockState.GetRuntime(&bestBlockHash)
 	if err != nil {
-		return err
+		return fmt.Errorf("getting runtime: %w", err)
 	}
 
-	rt, err := s.blockState.GetRuntime(&bestBlockHash)
+	stateVersion, err := rt.StateVersion()
 	if err != nil {
-		logger.Critical("failed to get runtime")
-		return err
+		return fmt.Errorf("getting runtime state version: %w", err)
+	}
+
+	ts, err := s.storageState.TrieState(stateRoot, stateVersion)
+	if err != nil {
+		return fmt.Errorf("computing trie state: %w", err)
 	}
 
 	rt.SetContextStorage(ts)
@@ -528,11 +547,8 @@ func (s *Service) HandleSubmittedExtrinsic(ext types.Extrinsic) error {
 }
 
 // GetMetadata calls runtime Metadata_metadata function
-func (s *Service) GetMetadata(bhash *common.Hash) ([]byte, error) {
-	var (
-		stateRootHash *common.Hash
-		err           error
-	)
+func (s *Service) GetMetadata(bhash *common.Hash) (metadata []byte, err error) {
+	var stateRootHash *common.Hash
 
 	// If block hash is not nil then fetch the state root corresponding to the block.
 	if bhash != nil {
@@ -541,12 +557,18 @@ func (s *Service) GetMetadata(bhash *common.Hash) ([]byte, error) {
 			return nil, err
 		}
 	}
-	ts, err := s.storageState.TrieState(stateRootHash)
+
+	rt, err := s.blockState.GetRuntime(bhash)
 	if err != nil {
 		return nil, err
 	}
 
-	rt, err := s.blockState.GetRuntime(bhash)
+	stateVersion, err := rt.StateVersion()
+	if err != nil {
+		return nil, fmt.Errorf("getting runtime state version: %w", err)
+	}
+
+	ts, err := s.storageState.TrieState(stateRootHash, stateVersion)
 	if err != nil {
 		return nil, err
 	}
@@ -563,12 +585,17 @@ func (s *Service) GetReadProofAt(block common.Hash, keys [][]byte) (
 		block = s.blockState.BestBlockHash()
 	}
 
+	stateVersion, err := s.blockState.GetRuntimeStateVersion(block)
+	if err != nil {
+		return hash, nil, fmt.Errorf("getting runtime state version: %w", err)
+	}
+
 	stateRoot, err := s.blockState.GetBlockStateRoot(block)
 	if err != nil {
 		return hash, nil, err
 	}
 
-	proofForKeys, err = s.storageState.GenerateTrieProof(stateRoot, keys)
+	proofForKeys, err = s.storageState.GenerateTrieProof(stateRoot, keys, stateVersion)
 	if err != nil {
 		return hash, nil, err
 	}

dot/import.go

@@ -13,6 +13,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/runtime/wasmer"
 	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 
@@ -21,11 +22,6 @@ 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)
-	if err != nil {
-		return err
-	}
-
 	header, err := newHeaderFromFile(headerFP)
 	if err != nil {
 		return err
@@ -37,11 +33,28 @@ func ImportState(basepath, stateFP, headerFP string, firstSlot uint64) error {
 		Path:     basepath,
 		LogLevel: log.Info,
 	}
-	srv := state.NewService(config)
-	return srv.Import(header, tr, firstSlot)
+	dotService := state.NewService(config)
+
+	blockHash := header.Hash()
+	runtimeCode, err := dotService.Storage.LoadCode(&blockHash)
+	if err != nil {
+		return fmt.Errorf("loading code from storage: %w", err)
+	}
+
+	stateVersion, err := wasmer.GetRuntimeStateVersion(runtimeCode)
+	if err != nil {
+		return fmt.Errorf("getting runtime state version: %w", err)
+	}
+
+	trie, err := newTrieFromPairs(stateFP, stateVersion)
+	if err != nil {
+		return fmt.Errorf("creating trie from pairs: %w", err)
+	}
+
+	return dotService.Import(header, trie, firstSlot, stateVersion)
 }
 
-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
@@ -63,7 +76,7 @@ func newTrieFromPairs(filename string) (*trie.Trie, error) {
 	}
 
 	tr := trie.NewEmptyTrie()
-	err = tr.LoadFromMap(entries)
+	err = tr.LoadFromMap(entries, stateVersion)
 	if err != nil {
 		return nil, err
 	}

dot/node.go

@@ -32,6 +32,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/grandpa"
 	"github.com/ChainSafe/gossamer/lib/keystore"
 	"github.com/ChainSafe/gossamer/lib/runtime"
+	"github.com/ChainSafe/gossamer/lib/runtime/wasmer"
 	"github.com/ChainSafe/gossamer/lib/services"
 	"github.com/ChainSafe/gossamer/lib/utils"
 )
@@ -143,14 +144,24 @@ func (*nodeBuilder) initNode(cfg *Config) error {
 		}
 	}
 
+	runtimeCode, err := gen.RuntimeCode()
+	if err != nil {
+		return fmt.Errorf("getting runtime code from genesis: %w", err)
+	}
+
+	stateVersion, err := wasmer.GetRuntimeStateVersion(runtimeCode)
+	if err != nil {
+		return fmt.Errorf("getting runtime state version: %w", err)
+	}
+
 	// create trie from genesis
-	t, err := genesis.NewTrieFromGenesis(gen)
+	t, err := genesis.NewTrieFromGenesis(gen, stateVersion)
 	if err != nil {
 		return fmt.Errorf("failed to create trie from genesis: %w", err)
 	}
 
 	// create genesis block from trie
-	header, err := genesis.NewGenesisBlockFromTrie(t)
+	header, err := genesis.NewGenesisBlockFromTrie(t, stateVersion)
 	if err != nil {
 		return fmt.Errorf("failed to create genesis block from trie: %w", err)
 	}