Inject tries to BlockState and StorageState

dot/rpc/modules/dev_integration_test.go

@@ -43,7 +43,8 @@ func newState(t *testing.T) (*state.BlockState, *state.EpochState) {
 	db := state.NewInMemoryDB(t)
 
 	_, genesisTrie, genesisHeader := genesis.NewTestGenesisWithTrieAndHeader(t)
-	bs, err := state.NewBlockStateFromGenesis(db, genesisTrie, genesisHeader, telemetryMock)
+	tries := state.NewTries(genesisTrie)
+	bs, err := state.NewBlockStateFromGenesis(db, tries, genesisHeader, telemetryMock)
 	require.NoError(t, err)
 	es, err := state.NewEpochStateFromGenesis(db, bs, genesisBABEConfig)
 	require.NoError(t, err)

dot/state/block.go

@@ -18,7 +18,6 @@ import (
 	"github.com/ChainSafe/gossamer/lib/blocktree"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/runtime"
-	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -60,7 +59,7 @@ type BlockState struct {
 	genesisHash       common.Hash
 	lastFinalised     common.Hash
 	unfinalisedBlocks *sync.Map // map[common.Hash]*types.Block
-	tries             *tries
+	tries             *Tries
 
 	// block notifiers
 	imported                       map[chan *types.Block]struct{}
@@ -74,7 +73,7 @@ type BlockState struct {
 }
 
 // NewBlockState will create a new BlockState backed by the database located at basePath
-func NewBlockState(db chaindb.Database, trs *tries, telemetry telemetry.Client) (*BlockState, error) {
+func NewBlockState(db chaindb.Database, trs *Tries, telemetry telemetry.Client) (*BlockState, error) {
 	bs := &BlockState{
 		dbPath:                     db.Path(),
 		baseState:                  NewBaseState(db),
@@ -106,14 +105,14 @@ func NewBlockState(db chaindb.Database, trs *tries, telemetry telemetry.Client)
 
 // NewBlockStateFromGenesis initialises a BlockState from a genesis header,
 // saving it to the database located at basePath
-func NewBlockStateFromGenesis(db chaindb.Database, trie *trie.Trie, header *types.Header,
-	telemetryMailer telemetry.Client) (*BlockState, error) {
+func NewBlockStateFromGenesis(db chaindb.Database, trs *Tries, header *types.Header,
+	telemetryMailer telemetry.Client) (*BlockState, error) { // TODO CHECKTEST
 	bs := &BlockState{
 		bt:                         blocktree.NewBlockTreeFromRoot(header),
 		baseState:                  NewBaseState(db),
 		db:                         chaindb.NewTable(db, blockPrefix),
 		unfinalisedBlocks:          new(sync.Map),
-		tries:                      newTries(trie),
+		tries:                      trs,
 		imported:                   make(map[chan *types.Block]struct{}),
 		finalised:                  make(map[chan *types.FinalisationInfo]struct{}),
 		runtimeUpdateSubscriptions: make(map[uint32]chan<- runtime.Version),

dot/state/block_data_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 func TestGetSet_ReceiptMessageQueue_Justification(t *testing.T) {
-	s := newTestBlockState(t, nil)
+	s := newTestBlockState(t, nil, newTriesEmpty())
 	require.NotNil(t, s)
 
 	var genesisHeader = &types.Header{

dot/state/block_finalisation.go

@@ -144,7 +144,7 @@ func (bs *BlockState) SetFinalisedHash(hash common.Hash, round, setID uint64) er
 		bs.notifyFinalized(hash, round, setID)
 	}
 
-	pruned := bs.bt.Prune(hash) // TODO this is always 0
+	pruned := bs.bt.Prune(hash)
 	for _, hash := range pruned {
 		block, has := bs.getAndDeleteUnfinalisedBlock(hash)
 		if !has {

dot/state/block_finalisation_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 func TestHighestRoundAndSetID(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	round, setID, err := bs.GetHighestRoundAndSetID()
 	require.NoError(t, err)
 	require.Equal(t, uint64(0), round)
@@ -61,7 +61,7 @@ func TestHighestRoundAndSetID(t *testing.T) {
 }
 
 func TestBlockState_SetFinalisedHash(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	h, err := bs.GetFinalisedHash(0, 0)
 	require.NoError(t, err)
 	require.Equal(t, testGenesisHeader.Hash(), h)
@@ -97,7 +97,7 @@ func TestBlockState_SetFinalisedHash(t *testing.T) {
 }
 
 func TestSetFinalisedHash_setFirstSlotOnFinalisation(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	firstSlot := uint64(42069)
 
 	digest := types.NewDigest()

dot/state/block_notify_test.go

@@ -12,13 +12,14 @@ import (
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	runtimemocks "github.com/ChainSafe/gossamer/lib/runtime/mocks"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/stretchr/testify/require"
 )
 
 var testMessageTimeout = time.Second * 3
 
 func TestImportChannel(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 	ch := bs.GetImportedBlockNotifierChannel()
 
 	defer bs.FreeImportedBlockNotifierChannel(ch)
@@ -35,7 +36,7 @@ func TestImportChannel(t *testing.T) {
 }
 
 func TestFreeImportedBlockNotifierChannel(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 	ch := bs.GetImportedBlockNotifierChannel()
 	require.Equal(t, 1, len(bs.imported))
 
@@ -44,7 +45,7 @@ func TestFreeImportedBlockNotifierChannel(t *testing.T) {
 }
 
 func TestFinalizedChannel(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 
 	ch := bs.GetFinalisedNotifierChannel()
 
@@ -66,7 +67,7 @@ func TestFinalizedChannel(t *testing.T) {
 }
 
 func TestImportChannel_Multi(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 
 	num := 5
 	chs := make([]chan *types.Block, num)
@@ -99,7 +100,7 @@ func TestImportChannel_Multi(t *testing.T) {
 }
 
 func TestFinalizedChannel_Multi(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 
 	num := 5
 	chs := make([]chan *types.FinalisationInfo, num)
@@ -136,7 +137,7 @@ func TestFinalizedChannel_Multi(t *testing.T) {
 }
 
 func TestService_RegisterUnRegisterRuntimeUpdatedChannel(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 	ch := make(chan<- runtime.Version)
 	chID, err := bs.RegisterRuntimeUpdatedChannel(ch)
 	require.NoError(t, err)
@@ -147,7 +148,7 @@ func TestService_RegisterUnRegisterRuntimeUpdatedChannel(t *testing.T) {
 }
 
 func TestService_RegisterUnRegisterConcurrentCalls(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 
 	go func() {
 		for i := 0; i < 100; i++ {

dot/state/block_race_test.go

@@ -28,15 +28,15 @@ func TestConcurrencySetHeader(t *testing.T) {
 		dbs[i] = NewInMemoryDB(t)
 	}
 
-	someTrie := trie.NewEmptyTrie() // not used in this test
+	tries := NewTries(trie.NewEmptyTrie()) // not used in this test
 
 	pend := new(sync.WaitGroup)
 	pend.Add(threads)
 	for i := 0; i < threads; i++ {
 		go func(index int) {
 			defer pend.Done()
 
-			bs, err := NewBlockStateFromGenesis(dbs[index], someTrie, testGenesisHeader, telemetryMock)
+			bs, err := NewBlockStateFromGenesis(dbs[index], tries, testGenesisHeader, telemetryMock)
 			require.NoError(t, err)
 
 			header := &types.Header{

dot/state/block_test.go

@@ -25,7 +25,7 @@ var testGenesisHeader = &types.Header{
 	Digest:    types.NewDigest(),
 }
 
-func newTestBlockState(t *testing.T, header *types.Header) *BlockState {
+func newTestBlockState(t *testing.T, header *types.Header, tries *Tries) *BlockState {
 	ctrl := gomock.NewController(t)
 	telemetryMock := NewMockClient(ctrl)
 	telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
@@ -35,13 +35,13 @@ func newTestBlockState(t *testing.T, header *types.Header) *BlockState {
 		header = testGenesisHeader
 	}
 
-	bs, err := NewBlockStateFromGenesis(db, trie.NewEmptyTrie(), header, telemetryMock)
+	bs, err := NewBlockStateFromGenesis(db, tries, header, telemetryMock)
 	require.NoError(t, err)
 	return bs
 }
 
 func TestSetAndGetHeader(t *testing.T) {
-	bs := newTestBlockState(t, nil)
+	bs := newTestBlockState(t, nil, newTriesEmpty())
 
 	header := &types.Header{
 		Number:    big.NewInt(0),
@@ -58,7 +58,7 @@ func TestSetAndGetHeader(t *testing.T) {
 }
 
 func TestHasHeader(t *testing.T) {
-	bs := newTestBlockState(t, nil)
+	bs := newTestBlockState(t, nil, newTriesEmpty())
 
 	header := &types.Header{
 		Number:    big.NewInt(0),
@@ -75,7 +75,7 @@ func TestHasHeader(t *testing.T) {
 }
 
 func TestGetBlockByNumber(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 
 	blockHeader := &types.Header{
 		ParentHash: testGenesisHeader.Hash(),
@@ -97,7 +97,7 @@ func TestGetBlockByNumber(t *testing.T) {
 }
 
 func TestAddBlock(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 
 	// Create header
 	header0 := &types.Header{
@@ -160,7 +160,7 @@ func TestAddBlock(t *testing.T) {
 }
 
 func TestGetSlotForBlock(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	expectedSlot := uint64(77)
 
 	babeHeader := types.NewBabeDigest()
@@ -191,7 +191,7 @@ func TestGetSlotForBlock(t *testing.T) {
 }
 
 func TestIsBlockOnCurrentChain(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	currChain, branchChains := AddBlocksToState(t, bs, 3, false)
 
 	for _, header := range currChain {
@@ -214,7 +214,7 @@ func TestIsBlockOnCurrentChain(t *testing.T) {
 }
 
 func TestAddBlock_BlockNumberToHash(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	currChain, branchChains := AddBlocksToState(t, bs, 8, false)
 
 	bestHash := bs.BestBlockHash()
@@ -262,7 +262,7 @@ func TestAddBlock_BlockNumberToHash(t *testing.T) {
 }
 
 func TestFinalization_DeleteBlock(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	AddBlocksToState(t, bs, 5, false)
 
 	btBefore := bs.bt.DeepCopy()
@@ -317,7 +317,7 @@ func TestFinalization_DeleteBlock(t *testing.T) {
 }
 
 func TestGetHashByNumber(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 
 	res, err := bs.GetHashByNumber(big.NewInt(0))
 	require.NoError(t, err)
@@ -344,7 +344,7 @@ func TestGetHashByNumber(t *testing.T) {
 
 func TestAddBlock_WithReOrg(t *testing.T) {
 	t.Skip() // TODO: this should be fixed after state refactor PR
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 
 	header1a := &types.Header{
 		Number:     big.NewInt(1),
@@ -453,7 +453,7 @@ func TestAddBlock_WithReOrg(t *testing.T) {
 }
 
 func TestAddBlockToBlockTree(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 
 	header := &types.Header{
 		Number:     big.NewInt(1),
@@ -473,7 +473,7 @@ func TestAddBlockToBlockTree(t *testing.T) {
 }
 
 func TestNumberIsFinalised(t *testing.T) {
-	bs := newTestBlockState(t, testGenesisHeader)
+	bs := newTestBlockState(t, testGenesisHeader, newTriesEmpty())
 	fin, err := bs.NumberIsFinalised(big.NewInt(0))
 	require.NoError(t, err)
 	require.True(t, fin)

dot/state/epoch_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
 	"github.com/ChainSafe/gossamer/lib/keystore"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 
 	"github.com/stretchr/testify/require"
@@ -28,7 +29,8 @@ var genesisBABEConfig = &types.BabeConfiguration{
 
 func newEpochStateFromGenesis(t *testing.T) *EpochState {
 	db := NewInMemoryDB(t)
-	s, err := NewEpochStateFromGenesis(db, newTestBlockState(t, nil), genesisBABEConfig)
+	blockState := newTestBlockState(t, nil, NewTries(trie.NewEmptyTrie()))
+	s, err := NewEpochStateFromGenesis(db, blockState, genesisBABEConfig)
 	require.NoError(t, err)
 	return s
 }
@@ -184,7 +186,7 @@ func TestEpochState_SetAndGetSlotDuration(t *testing.T) {
 
 func TestEpochState_GetEpochFromTime(t *testing.T) {
 	s := newEpochStateFromGenesis(t)
-	s.blockState = newTestBlockState(t, testGenesisHeader)
+	s.blockState = newTestBlockState(t, testGenesisHeader, NewTries(trie.NewEmptyTrie()))
 
 	epochDuration, err := time.ParseDuration(
 		fmt.Sprintf("%dms",

dot/state/helpers_test.go

@@ -8,9 +8,17 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/stretchr/testify/require"
 )
 
+func newTriesEmpty() *Tries {
+	return &Tries{
+		rootToTrie: make(map[common.Hash]*trie.Trie),
+	}
+}
+
 // newGenerator creates a new PRNG seeded with the
 // unix nanoseconds value of the current time.
 func newGenerator() (prng *rand.Rand) {

dot/state/initialize.go

@@ -62,14 +62,16 @@ func (s *Service) Initialise(gen *genesis.Genesis, header *types.Header, t *trie
 		return fmt.Errorf("failed to write genesis values to database: %s", err)
 	}
 
+	tries := NewTries(t)
+
 	// create block state from genesis block
-	blockState, err := NewBlockStateFromGenesis(db, t, header, s.Telemetry)
+	blockState, err := NewBlockStateFromGenesis(db, tries, header, s.Telemetry)
 	if err != nil {
 		return fmt.Errorf("failed to create block state from genesis: %s", err)
 	}
 
 	// create storage state from genesis trie
-	storageState, err := NewStorageState(db, blockState, pruner.Config{})
+	storageState, err := NewStorageState(db, blockState, tries, pruner.Config{})
 	if err != nil {
 		return fmt.Errorf("failed to create storage state from trie: %s", err)
 	}

dot/state/offline_pruner.go

@@ -41,7 +41,7 @@ func NewOfflinePruner(inputDBPath, prunedDBPath string, bloomSize uint64,
 		return nil, fmt.Errorf("failed to load DB %w", err)
 	}
 
-	tries := newTries(trie.NewEmptyTrie())
+	tries := NewTries(trie.NewEmptyTrie())
 
 	// create blockState state
 	// NewBlockState on pruner execution does not use telemetry
@@ -62,7 +62,7 @@ func NewOfflinePruner(inputDBPath, prunedDBPath string, bloomSize uint64,
 	}
 
 	// load storage state
-	storageState, err := NewStorageState(db, blockState, pruner.Config{})
+	storageState, err := NewStorageState(db, blockState, tries, pruner.Config{})
 	if err != nil {
 		return nil, fmt.Errorf("failed to create new storage state %w", err)
 	}