all: introduce trie owner notion (#24750)

* cmd, core/state, light, trie, eth: add trie owner notion

* all: refactor

* tests: fix goimports

* core/state/snapshot: fix ineffasigns

Co-authored-by: Martin Holst Swende <martin@swende.se>

core/blockchain.go

@@ -674,7 +674,7 @@ func (bc *BlockChain) SnapSyncCommitHead(hash common.Hash) error {
 	if block == nil {
 		return fmt.Errorf("non existent block [%x..]", hash[:4])
 	}
-	if _, err := trie.NewSecure(block.Root(), bc.stateCache.TrieDB()); err != nil {
+	if _, err := trie.NewSecure(common.Hash{}, block.Root(), bc.stateCache.TrieDB()); err != nil {
 		return err
 	}
 

core/state/database.go

@@ -133,7 +133,7 @@ type cachingDB struct {
 
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
-	tr, err := trie.NewSecure(root, db.db)
+	tr, err := trie.NewSecure(common.Hash{}, root, db.db)
 	if err != nil {
 		return nil, err
 	}
@@ -142,7 +142,7 @@ func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
 
 // OpenStorageTrie opens the storage trie of an account.
 func (db *cachingDB) OpenStorageTrie(addrHash, root common.Hash) (Trie, error) {
-	tr, err := trie.NewSecure(root, db.db)
+	tr, err := trie.NewSecure(addrHash, root, db.db)
 	if err != nil {
 		return nil, err
 	}

core/state/pruner/pruner.go

@@ -410,7 +410,7 @@ func extractGenesis(db ethdb.Database, stateBloom *stateBloom) error {
 	if genesis == nil {
 		return errors.New("missing genesis block")
 	}
-	t, err := trie.NewSecure(genesis.Root(), trie.NewDatabase(db))
+	t, err := trie.NewSecure(common.Hash{}, genesis.Root(), trie.NewDatabase(db))
 	if err != nil {
 		return err
 	}
@@ -430,7 +430,7 @@ func extractGenesis(db ethdb.Database, stateBloom *stateBloom) error {
 				return err
 			}
 			if acc.Root != emptyRoot {
-				storageTrie, err := trie.NewSecure(acc.Root, trie.NewDatabase(db))
+				storageTrie, err := trie.NewSecure(common.BytesToHash(accIter.LeafKey()), acc.Root, trie.NewDatabase(db))
 				if err != nil {
 					return err
 				}

core/state/snapshot/conversion.go

@@ -43,7 +43,7 @@ type trieKV struct {
 type (
 	// trieGeneratorFn is the interface of trie generation which can
 	// be implemented by different trie algorithm.
-	trieGeneratorFn func(db ethdb.KeyValueWriter, in chan (trieKV), out chan (common.Hash))
+	trieGeneratorFn func(db ethdb.KeyValueWriter, owner common.Hash, in chan (trieKV), out chan (common.Hash))
 
 	// leafCallbackFn is the callback invoked at the leaves of the trie,
 	// returns the subtrie root with the specified subtrie identifier.
@@ -253,7 +253,7 @@ func generateTrieRoot(db ethdb.KeyValueWriter, it Iterator, account common.Hash,
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		generatorFn(db, in, out)
+		generatorFn(db, account, in, out)
 	}()
 	// Spin up a go-routine for progress logging
 	if report && stats != nil {
@@ -360,8 +360,8 @@ func generateTrieRoot(db ethdb.KeyValueWriter, it Iterator, account common.Hash,
 	return stop(nil)
 }
 
-func stackTrieGenerate(db ethdb.KeyValueWriter, in chan trieKV, out chan common.Hash) {
-	t := trie.NewStackTrie(db)
+func stackTrieGenerate(db ethdb.KeyValueWriter, owner common.Hash, in chan trieKV, out chan common.Hash) {
+	t := trie.NewStackTrieWithOwner(db, owner)
 	for leaf := range in {
 		t.TryUpdate(leaf.key[:], leaf.value)
 	}

core/state/snapshot/generate.go

@@ -166,7 +166,7 @@ func (result *proofResult) forEach(callback func(key []byte, val []byte) error)
 //
 // The proof result will be returned if the range proving is finished, otherwise
 // the error will be returned to abort the entire procedure.
-func (dl *diskLayer) proveRange(ctx *generatorContext, root common.Hash, prefix []byte, kind string, origin []byte, max int, valueConvertFn func([]byte) ([]byte, error)) (*proofResult, error) {
+func (dl *diskLayer) proveRange(ctx *generatorContext, owner common.Hash, root common.Hash, prefix []byte, kind string, origin []byte, max int, valueConvertFn func([]byte) ([]byte, error)) (*proofResult, error) {
 	var (
 		keys     [][]byte
 		vals     [][]byte
@@ -234,7 +234,7 @@ func (dl *diskLayer) proveRange(ctx *generatorContext, root common.Hash, prefix
 
 	// The snap state is exhausted, pass the entire key/val set for verification
 	if origin == nil && !diskMore {
-		stackTr := trie.NewStackTrie(nil)
+		stackTr := trie.NewStackTrieWithOwner(nil, owner)
 		for i, key := range keys {
 			stackTr.TryUpdate(key, vals[i])
 		}
@@ -248,7 +248,7 @@ func (dl *diskLayer) proveRange(ctx *generatorContext, root common.Hash, prefix
 		return &proofResult{keys: keys, vals: vals}, nil
 	}
 	// Snap state is chunked, generate edge proofs for verification.
-	tr, err := trie.New(root, dl.triedb)
+	tr, err := trie.New(owner, root, dl.triedb)
 	if err != nil {
 		ctx.stats.Log("Trie missing, state snapshotting paused", dl.root, dl.genMarker)
 		return nil, errMissingTrie
@@ -313,9 +313,9 @@ type onStateCallback func(key []byte, val []byte, write bool, delete bool) error
 // generateRange generates the state segment with particular prefix. Generation can
 // either verify the correctness of existing state through range-proof and skip
 // generation, or iterate trie to regenerate state on demand.
-func (dl *diskLayer) generateRange(ctx *generatorContext, root common.Hash, prefix []byte, kind string, origin []byte, max int, onState onStateCallback, valueConvertFn func([]byte) ([]byte, error)) (bool, []byte, error) {
+func (dl *diskLayer) generateRange(ctx *generatorContext, owner common.Hash, root common.Hash, prefix []byte, kind string, origin []byte, max int, onState onStateCallback, valueConvertFn func([]byte) ([]byte, error)) (bool, []byte, error) {
 	// Use range prover to check the validity of the flat state in the range
-	result, err := dl.proveRange(ctx, root, prefix, kind, origin, max, valueConvertFn)
+	result, err := dl.proveRange(ctx, owner, root, prefix, kind, origin, max, valueConvertFn)
 	if err != nil {
 		return false, nil, err
 	}
@@ -363,7 +363,7 @@ func (dl *diskLayer) generateRange(ctx *generatorContext, root common.Hash, pref
 	if len(result.keys) > 0 {
 		snapNodeCache = memorydb.New()
 		snapTrieDb := trie.NewDatabase(snapNodeCache)
-		snapTrie, _ := trie.New(common.Hash{}, snapTrieDb)
+		snapTrie, _ := trie.New(owner, common.Hash{}, snapTrieDb)
 		for i, key := range result.keys {
 			snapTrie.Update(key, result.vals[i])
 		}
@@ -374,7 +374,7 @@ func (dl *diskLayer) generateRange(ctx *generatorContext, root common.Hash, pref
 	// if it's already opened with some nodes resolved.
 	tr := result.tr
 	if tr == nil {
-		tr, err = trie.New(root, dl.triedb)
+		tr, err = trie.New(owner, root, dl.triedb)
 		if err != nil {
 			ctx.stats.Log("Trie missing, state snapshotting paused", dl.root, dl.genMarker)
 			return false, nil, errMissingTrie
@@ -537,7 +537,7 @@ func generateStorages(ctx *generatorContext, dl *diskLayer, account common.Hash,
 	// Loop for re-generating the missing storage slots.
 	var origin = common.CopyBytes(storeMarker)
 	for {
-		exhausted, last, err := dl.generateRange(ctx, storageRoot, append(rawdb.SnapshotStoragePrefix, account.Bytes()...), snapStorage, origin, storageCheckRange, onStorage, nil)
+		exhausted, last, err := dl.generateRange(ctx, account, storageRoot, append(rawdb.SnapshotStoragePrefix, account.Bytes()...), snapStorage, origin, storageCheckRange, onStorage, nil)
 		if err != nil {
 			return err // The procedure it aborted, either by external signal or internal error.
 		}
@@ -637,7 +637,7 @@ func generateAccounts(ctx *generatorContext, dl *diskLayer, accMarker []byte) er
 	}
 	origin := common.CopyBytes(accMarker)
 	for {
-		exhausted, last, err := dl.generateRange(ctx, dl.root, rawdb.SnapshotAccountPrefix, snapAccount, origin, accountRange, onAccount, FullAccountRLP)
+		exhausted, last, err := dl.generateRange(ctx, common.Hash{}, dl.root, rawdb.SnapshotAccountPrefix, snapAccount, origin, accountRange, onAccount, FullAccountRLP)
 		if err != nil {
 			return err // The procedure it aborted, either by external signal or internal error.
 		}

core/state/snapshot/generate_test.go

@@ -26,61 +26,12 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/ethdb"
-	"github.com/ethereum/go-ethereum/ethdb/memorydb"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 	"golang.org/x/crypto/sha3"
 )
 
-// Tests that snapshot generation from an empty database.
-func TestGeneration(t *testing.T) {
-	// We can't use statedb to make a test trie (circular dependency), so make
-	// a fake one manually. We're going with a small account trie of 3 accounts,
-	// two of which also has the same 3-slot storage trie attached.
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	stTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	stTrie.Update([]byte("key-1"), []byte("val-1")) // 0x1314700b81afc49f94db3623ef1df38f3ed18b73a1b7ea2f6c095118cf6118a0
-	stTrie.Update([]byte("key-2"), []byte("val-2")) // 0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371
-	stTrie.Update([]byte("key-3"), []byte("val-3")) // 0x51c71a47af0695957647fb68766d0becee77e953df17c29b3c2f25436f055c78
-	stTrie.Commit(nil)                              // Root: 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
-
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ := rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
-
-	acc = &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-2"), val) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
-
-	acc = &Account{Balance: big.NewInt(3), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-3"), val) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
-	root, _, _ := accTrie.Commit(nil)    // Root: 0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd
-	triedb.Commit(root, false, nil)
-
-	if have, want := root, common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"); have != want {
-		t.Fatalf("have %#x want %#x", have, want)
-	}
-	snap := generateSnapshot(diskdb, triedb, 16, root)
-	select {
-	case <-snap.genPending:
-		// Snapshot generation succeeded
-
-	case <-time.After(3 * time.Second):
-		t.Errorf("Snapshot generation failed")
-	}
-	checkSnapRoot(t, snap, root)
-	// Signal abortion to the generator and wait for it to tear down
-	stop := make(chan *generatorStats)
-	snap.genAbort <- stop
-	<-stop
-}
-
 func hashData(input []byte) common.Hash {
 	var hasher = sha3.NewLegacyKeccak256()
 	var hash common.Hash
@@ -90,48 +41,25 @@ func hashData(input []byte) common.Hash {
 	return hash
 }
 
-// Tests that snapshot generation with existent flat state.
-func TestGenerateExistentState(t *testing.T) {
+// Tests that snapshot generation from an empty database.
+func TestGeneration(t *testing.T) {
 	// We can't use statedb to make a test trie (circular dependency), so make
 	// a fake one manually. We're going with a small account trie of 3 accounts,
 	// two of which also has the same 3-slot storage trie attached.
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	stTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	stTrie.Update([]byte("key-1"), []byte("val-1")) // 0x1314700b81afc49f94db3623ef1df38f3ed18b73a1b7ea2f6c095118cf6118a0
-	stTrie.Update([]byte("key-2"), []byte("val-2")) // 0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371
-	stTrie.Update([]byte("key-3"), []byte("val-3")) // 0x51c71a47af0695957647fb68766d0becee77e953df17c29b3c2f25436f055c78
-	stTrie.Commit(nil)                              // Root: 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
+	var helper = newHelper()
+	stRoot := helper.makeStorageTrie(common.Hash{}, common.Hash{}, []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, false)
 
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ := rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
-	rawdb.WriteAccountSnapshot(diskdb, hashData([]byte("acc-1")), val)
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-1")), hashData([]byte("key-1")), []byte("val-1"))
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-1")), hashData([]byte("key-2")), []byte("val-2"))
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-1")), hashData([]byte("key-3")), []byte("val-3"))
+	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
+	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
+	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()})
 
-	acc = &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-2"), val) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
-	diskdb.Put(hashData([]byte("acc-2")).Bytes(), val)
-	rawdb.WriteAccountSnapshot(diskdb, hashData([]byte("acc-2")), val)
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
-	acc = &Account{Balance: big.NewInt(3), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-3"), val) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
-	rawdb.WriteAccountSnapshot(diskdb, hashData([]byte("acc-3")), val)
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-3")), hashData([]byte("key-1")), []byte("val-1"))
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-3")), hashData([]byte("key-2")), []byte("val-2"))
-	rawdb.WriteStorageSnapshot(diskdb, hashData([]byte("acc-3")), hashData([]byte("key-3")), []byte("val-3"))
-
-	root, _, _ := accTrie.Commit(nil) // Root: 0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd
-	triedb.Commit(root, false, nil)
-
-	snap := generateSnapshot(diskdb, triedb, 16, root)
+	root, snap := helper.CommitAndGenerate()
+	if have, want := root, common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"); have != want {
+		t.Fatalf("have %#x want %#x", have, want)
+	}
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -140,6 +68,43 @@ func TestGenerateExistentState(t *testing.T) {
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
+	// Signal abortion to the generator and wait for it to tear down
+	stop := make(chan *generatorStats)
+	snap.genAbort <- stop
+	<-stop
+}
+
+// Tests that snapshot generation with existent flat state.
+func TestGenerateExistentState(t *testing.T) {
+	// We can't use statedb to make a test trie (circular dependency), so make
+	// a fake one manually. We're going with a small account trie of 3 accounts,
+	// two of which also has the same 3-slot storage trie attached.
+	var helper = newHelper()
+
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
+	helper.addSnapAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
+	helper.addSnapStorage("acc-1", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
+
+	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
+	helper.addSnapAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
+
+	stRoot = helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()})
+	helper.addSnapAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()})
+	helper.addSnapStorage("acc-3", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
+
+	root, snap := helper.CommitAndGenerate()
+	select {
+	case <-snap.genPending:
+		// Snapshot generation succeeded
+
+	case <-time.After(3 * time.Second):
+		t.Errorf("Snapshot generation failed")
+	}
+	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
 	stop := make(chan *generatorStats)
 	snap.genAbort <- stop
@@ -163,7 +128,6 @@ func checkSnapRoot(t *testing.T, snap *diskLayer, trieRoot common.Hash) {
 			}
 			return hash, nil
 		}, newGenerateStats(), true)
-
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -176,15 +140,15 @@ func checkSnapRoot(t *testing.T, snap *diskLayer, trieRoot common.Hash) {
 }
 
 type testHelper struct {
-	diskdb  *memorydb.Database
+	diskdb  ethdb.Database
 	triedb  *trie.Database
 	accTrie *trie.SecureTrie
 }
 
 func newHelper() *testHelper {
-	diskdb := memorydb.New()
+	diskdb := rawdb.NewMemoryDatabase()
 	triedb := trie.NewDatabase(diskdb)
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
+	accTrie, _ := trie.NewSecure(common.Hash{}, common.Hash{}, triedb)
 	return &testHelper{
 		diskdb:  diskdb,
 		triedb:  triedb,
@@ -215,18 +179,28 @@ func (t *testHelper) addSnapStorage(accKey string, keys []string, vals []string)
 	}
 }
 
-func (t *testHelper) makeStorageTrie(keys []string, vals []string) []byte {
-	stTrie, _ := trie.NewSecure(common.Hash{}, t.triedb)
+func (t *testHelper) makeStorageTrie(stateRoot, owner common.Hash, keys []string, vals []string, commit bool) []byte {
+	stTrie, _ := trie.NewSecure(owner, common.Hash{}, t.triedb)
 	for i, k := range keys {
 		stTrie.Update([]byte(k), []byte(vals[i]))
 	}
-	root, _, _ := stTrie.Commit(nil)
+	var root common.Hash
+	if !commit {
+		root = stTrie.Hash()
+	} else {
+		root, _, _ = stTrie.Commit(nil)
+	}
 	return root.Bytes()
 }
 
-func (t *testHelper) Generate() (common.Hash, *diskLayer) {
+func (t *testHelper) Commit() common.Hash {
 	root, _, _ := t.accTrie.Commit(nil)
 	t.triedb.Commit(root, false, nil)
+	return root
+}
+
+func (t *testHelper) CommitAndGenerate() (common.Hash, *diskLayer) {
+	root := t.Commit()
 	snap := generateSnapshot(t.diskdb, t.triedb, 16, root)
 	return root, snap
 }
@@ -249,26 +223,29 @@ func (t *testHelper) Generate() (common.Hash, *diskLayer) {
 //   - extra slots in the end
 func TestGenerateExistentStateWithWrongStorage(t *testing.T) {
 	helper := newHelper()
-	stRoot := helper.makeStorageTrie([]string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
 
 	// Account one, empty root but non-empty database
 	helper.addAccount("acc-1", &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
 	helper.addSnapStorage("acc-1", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
 
 	// Account two, non empty root but empty database
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-2")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 	helper.addAccount("acc-2", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 
 	// Miss slots
 	{
 		// Account three, non empty root but misses slots in the beginning
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-3", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-3", []string{"key-2", "key-3"}, []string{"val-2", "val-3"})
 
 		// Account four, non empty root but misses slots in the middle
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-4")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-4", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-4", []string{"key-1", "key-3"}, []string{"val-1", "val-3"})
 
 		// Account five, non empty root but misses slots in the end
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-5")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-5", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-5", []string{"key-1", "key-2"}, []string{"val-1", "val-2"})
 	}
@@ -276,18 +253,22 @@ func TestGenerateExistentStateWithWrongStorage(t *testing.T) {
 	// Wrong storage slots
 	{
 		// Account six, non empty root but wrong slots in the beginning
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-6")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-6", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-6", []string{"key-1", "key-2", "key-3"}, []string{"badval-1", "val-2", "val-3"})
 
 		// Account seven, non empty root but wrong slots in the middle
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-7")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-7", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-7", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "badval-2", "val-3"})
 
 		// Account eight, non empty root but wrong slots in the end
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-8")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-8", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-8", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "badval-3"})
 
 		// Account 9, non empty root but rotated slots
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-9")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-9", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-9", []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-3", "val-2"})
 	}
@@ -295,19 +276,22 @@ func TestGenerateExistentStateWithWrongStorage(t *testing.T) {
 	// Extra storage slots
 	{
 		// Account 10, non empty root but extra slots in the beginning
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-10")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-10", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-10", []string{"key-0", "key-1", "key-2", "key-3"}, []string{"val-0", "val-1", "val-2", "val-3"})
 
 		// Account 11, non empty root but extra slots in the middle
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-11")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-11", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-11", []string{"key-1", "key-2", "key-2-1", "key-3"}, []string{"val-1", "val-2", "val-2-1", "val-3"})
 
 		// Account 12, non empty root but extra slots in the end
+		helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-12")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addAccount("acc-12", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		helper.addSnapStorage("acc-12", []string{"key-1", "key-2", "key-3", "key-4"}, []string{"val-1", "val-2", "val-3", "val-4"})
 	}
 
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	t.Logf("Root: %#x\n", root) // Root = 0x8746cce9fd9c658b2cfd639878ed6584b7a2b3e73bb40f607fcfa156002429a0
 
 	select {
@@ -331,7 +315,12 @@ func TestGenerateExistentStateWithWrongStorage(t *testing.T) {
 // - extra accounts
 func TestGenerateExistentStateWithWrongAccounts(t *testing.T) {
 	helper := newHelper()
-	stRoot := helper.makeStorageTrie([]string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
+
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-2")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-4")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-6")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
 	// Trie accounts [acc-1, acc-2, acc-3, acc-4, acc-6]
 	// Extra accounts [acc-0, acc-5, acc-7]
@@ -359,7 +348,7 @@ func TestGenerateExistentStateWithWrongAccounts(t *testing.T) {
 		helper.addSnapAccount("acc-7", &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyRoot.Bytes()})          // after the end
 	}
 
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	t.Logf("Root: %#x\n", root) // Root = 0x825891472281463511e7ebcc7f109e4f9200c20fa384754e11fd605cd98464e8
 
 	select {
@@ -383,29 +372,19 @@ func TestGenerateCorruptAccountTrie(t *testing.T) {
 	// We can't use statedb to make a test trie (circular dependency), so make
 	// a fake one manually. We're going with a small account trie of 3 accounts,
 	// without any storage slots to keep the test smaller.
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	tr, _ := trie.NewSecure(common.Hash{}, triedb)
-	acc := &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ := rlp.EncodeToBytes(acc)
-	tr.Update([]byte("acc-1"), val) // 0xc7a30f39aff471c95d8a837497ad0e49b65be475cc0953540f80cfcdbdcd9074
+	helper := newHelper()
 
-	acc = &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	tr.Update([]byte("acc-2"), val) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
+	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}) // 0xc7a30f39aff471c95d8a837497ad0e49b65be475cc0953540f80cfcdbdcd9074
+	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
+	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}) // 0x19ead688e907b0fab07176120dceec244a72aff2f0aa51e8b827584e378772f4
 
-	acc = &Account{Balance: big.NewInt(3), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	tr.Update([]byte("acc-3"), val) // 0x19ead688e907b0fab07176120dceec244a72aff2f0aa51e8b827584e378772f4
-	tr.Commit(nil)                  // Root: 0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978
+	root, _, _ := helper.accTrie.Commit(nil) // Root: 0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978
 
 	// Delete an account trie leaf and ensure the generator chokes
-	triedb.Commit(common.HexToHash("0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978"), false, nil)
-	diskdb.Delete(common.HexToHash("0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7").Bytes())
+	helper.triedb.Commit(root, false, nil)
+	helper.diskdb.Delete(common.HexToHash("0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7").Bytes())
 
-	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978"))
+	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -427,45 +406,30 @@ func TestGenerateMissingStorageTrie(t *testing.T) {
 	// We can't use statedb to make a test trie (circular dependency), so make
 	// a fake one manually. We're going with a small account trie of 3 accounts,
 	// two of which also has the same 3-slot storage trie attached.
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	stTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	stTrie.Update([]byte("key-1"), []byte("val-1")) // 0x1314700b81afc49f94db3623ef1df38f3ed18b73a1b7ea2f6c095118cf6118a0
-	stTrie.Update([]byte("key-2"), []byte("val-2")) // 0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371
-	stTrie.Update([]byte("key-3"), []byte("val-3")) // 0x51c71a47af0695957647fb68766d0becee77e953df17c29b3c2f25436f055c78
-	stTrie.Commit(nil)                              // Root: 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
+	helper := newHelper()
 
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ := rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
-
-	acc = &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-2"), val) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
-
-	acc = &Account{Balance: big.NewInt(3), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-3"), val) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
-	accTrie.Commit(nil)                  // Root: 0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true) // 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
+	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})                                                // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})                                     // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
+	stRoot = helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()}) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
+	root, _, _ := helper.accTrie.Commit(nil)
 
 	// We can only corrupt the disk database, so flush the tries out
-	triedb.Reference(
-		common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67"),
+	helper.triedb.Reference(
+		common.BytesToHash(stRoot),
 		common.HexToHash("0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e"),
 	)
-	triedb.Reference(
-		common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67"),
+	helper.triedb.Reference(
+		common.BytesToHash(stRoot),
 		common.HexToHash("0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2"),
 	)
-	triedb.Commit(common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"), false, nil)
+	helper.triedb.Commit(root, false, nil)
 
 	// Delete a storage trie root and ensure the generator chokes
-	diskdb.Delete(common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67").Bytes())
+	helper.diskdb.Delete(stRoot)
 
-	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"))
+	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -486,45 +450,31 @@ func TestGenerateCorruptStorageTrie(t *testing.T) {
 	// We can't use statedb to make a test trie (circular dependency), so make
 	// a fake one manually. We're going with a small account trie of 3 accounts,
 	// two of which also has the same 3-slot storage trie attached.
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	stTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	stTrie.Update([]byte("key-1"), []byte("val-1")) // 0x1314700b81afc49f94db3623ef1df38f3ed18b73a1b7ea2f6c095118cf6118a0
-	stTrie.Update([]byte("key-2"), []byte("val-2")) // 0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371
-	stTrie.Update([]byte("key-3"), []byte("val-3")) // 0x51c71a47af0695957647fb68766d0becee77e953df17c29b3c2f25436f055c78
-	stTrie.Commit(nil)                              // Root: 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
+	helper := newHelper()
 
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ := rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+	stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true) // 0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67
+	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})                                                // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})                                     // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
+	stRoot = helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-3")), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
+	helper.addTrieAccount("acc-3", &Account{Balance: big.NewInt(3), Root: stRoot, CodeHash: emptyCode.Bytes()}) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
 
-	acc = &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-2"), val) // 0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7
-
-	acc = &Account{Balance: big.NewInt(3), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
-	val, _ = rlp.EncodeToBytes(acc)
-	accTrie.Update([]byte("acc-3"), val) // 0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2
-	accTrie.Commit(nil)                  // Root: 0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd
+	root, _, _ := helper.accTrie.Commit(nil)
 
 	// We can only corrupt the disk database, so flush the tries out
-	triedb.Reference(
-		common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67"),
+	helper.triedb.Reference(
+		common.BytesToHash(stRoot),
 		common.HexToHash("0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e"),
 	)
-	triedb.Reference(
-		common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67"),
+	helper.triedb.Reference(
+		common.BytesToHash(stRoot),
 		common.HexToHash("0x50815097425d000edfc8b3a4a13e175fc2bdcfee8bdfbf2d1ff61041d3c235b2"),
 	)
-	triedb.Commit(common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"), false, nil)
+	helper.triedb.Commit(root, false, nil)
 
 	// Delete a storage trie leaf and ensure the generator chokes
-	diskdb.Delete(common.HexToHash("0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371").Bytes())
+	helper.diskdb.Delete(common.HexToHash("0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371").Bytes())
 
-	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0xe3712f1a226f3782caca78ca770ccc19ee000552813a9f59d479f8611db9b1fd"))
+	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -539,56 +489,51 @@ func TestGenerateCorruptStorageTrie(t *testing.T) {
 	<-stop
 }
 
-func getStorageTrie(n int, triedb *trie.Database) *trie.SecureTrie {
-	stTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	for i := 0; i < n; i++ {
-		k := fmt.Sprintf("key-%d", i)
-		v := fmt.Sprintf("val-%d", i)
-		stTrie.Update([]byte(k), []byte(v))
-	}
-	stTrie.Commit(nil)
-	return stTrie
-}
-
 // Tests that snapshot generation when an extra account with storage exists in the snap state.
 func TestGenerateWithExtraAccounts(t *testing.T) {
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-		stTrie = getStorageTrie(5, triedb)
-	)
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	{ // Account one in the trie
-		acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
+	helper := newHelper()
+	{
+		// Account one in the trie
+		stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")),
+			[]string{"key-1", "key-2", "key-3", "key-4", "key-5"},
+			[]string{"val-1", "val-2", "val-3", "val-4", "val-5"},
+			true,
+		)
+		acc := &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()}
 		val, _ := rlp.EncodeToBytes(acc)
-		accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+		helper.accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+
 		// Identical in the snap
 		key := hashData([]byte("acc-1"))
-		rawdb.WriteAccountSnapshot(diskdb, key, val)
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-1")), []byte("val-1"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-2")), []byte("val-2"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-3")), []byte("val-3"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-4")), []byte("val-4"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-5")), []byte("val-5"))
+		rawdb.WriteAccountSnapshot(helper.triedb.DiskDB(), key, val)
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("key-1")), []byte("val-1"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("key-2")), []byte("val-2"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("key-3")), []byte("val-3"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("key-4")), []byte("val-4"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("key-5")), []byte("val-5"))
 	}
-	{ // Account two exists only in the snapshot
-		acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
+	{
+		// Account two exists only in the snapshot
+		stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-2")),
+			[]string{"key-1", "key-2", "key-3", "key-4", "key-5"},
+			[]string{"val-1", "val-2", "val-3", "val-4", "val-5"},
+			true,
+		)
+		acc := &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()}
 		val, _ := rlp.EncodeToBytes(acc)
 		key := hashData([]byte("acc-2"))
-		rawdb.WriteAccountSnapshot(diskdb, key, val)
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("b-key-1")), []byte("b-val-1"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("b-key-2")), []byte("b-val-2"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("b-key-3")), []byte("b-val-3"))
+		rawdb.WriteAccountSnapshot(helper.triedb.DiskDB(), key, val)
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("b-key-1")), []byte("b-val-1"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("b-key-2")), []byte("b-val-2"))
+		rawdb.WriteStorageSnapshot(helper.triedb.DiskDB(), key, hashData([]byte("b-key-3")), []byte("b-val-3"))
 	}
-	root, _, _ := accTrie.Commit(nil)
-	t.Logf("root: %x", root)
-	triedb.Commit(root, false, nil)
+	root := helper.Commit()
+
 	// To verify the test: If we now inspect the snap db, there should exist extraneous storage items
-	if data := rawdb.ReadStorageSnapshot(diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data == nil {
+	if data := rawdb.ReadStorageSnapshot(helper.triedb.DiskDB(), hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data == nil {
 		t.Fatalf("expected snap storage to exist")
 	}
-
-	snap := generateSnapshot(diskdb, triedb, 16, root)
+	snap := generateSnapshot(helper.diskdb, helper.triedb, 16, root)
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -597,12 +542,13 @@ func TestGenerateWithExtraAccounts(t *testing.T) {
 		t.Errorf("Snapshot generation failed")
 	}
 	checkSnapRoot(t, snap, root)
+
 	// Signal abortion to the generator and wait for it to tear down
 	stop := make(chan *generatorStats)
 	snap.genAbort <- stop
 	<-stop
 	// If we now inspect the snap db, there should exist no extraneous storage items
-	if data := rawdb.ReadStorageSnapshot(diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
+	if data := rawdb.ReadStorageSnapshot(helper.triedb.DiskDB(), hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
 		t.Fatalf("expected slot to be removed, got %v", string(data))
 	}
 }
@@ -616,37 +562,36 @@ func TestGenerateWithManyExtraAccounts(t *testing.T) {
 	if false {
 		enableLogging()
 	}
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-		stTrie = getStorageTrie(3, triedb)
-	)
-	accTrie, _ := trie.NewSecure(common.Hash{}, triedb)
-	{ // Account one in the trie
-		acc := &Account{Balance: big.NewInt(1), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
+	helper := newHelper()
+	{
+		// Account one in the trie
+		stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte("acc-1")),
+			[]string{"key-1", "key-2", "key-3"},
+			[]string{"val-1", "val-2", "val-3"},
+			true,
+		)
+		acc := &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()}
 		val, _ := rlp.EncodeToBytes(acc)
-		accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+		helper.accTrie.Update([]byte("acc-1"), val) // 0x9250573b9c18c664139f3b6a7a8081b7d8f8916a8fcc5d94feec6c29f5fd4e9e
+
 		// Identical in the snap
 		key := hashData([]byte("acc-1"))
-		rawdb.WriteAccountSnapshot(diskdb, key, val)
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-1")), []byte("val-1"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-2")), []byte("val-2"))
-		rawdb.WriteStorageSnapshot(diskdb, key, hashData([]byte("key-3")), []byte("val-3"))
+		rawdb.WriteAccountSnapshot(helper.diskdb, key, val)
+		rawdb.WriteStorageSnapshot(helper.diskdb, key, hashData([]byte("key-1")), []byte("val-1"))
+		rawdb.WriteStorageSnapshot(helper.diskdb, key, hashData([]byte("key-2")), []byte("val-2"))
+		rawdb.WriteStorageSnapshot(helper.diskdb, key, hashData([]byte("key-3")), []byte("val-3"))
 	}
-	{ // 100 accounts exist only in snapshot
+	{
+		// 100 accounts exist only in snapshot
 		for i := 0; i < 1000; i++ {
 			//acc := &Account{Balance: big.NewInt(int64(i)), Root: stTrie.Hash().Bytes(), CodeHash: emptyCode.Bytes()}
 			acc := &Account{Balance: big.NewInt(int64(i)), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
 			val, _ := rlp.EncodeToBytes(acc)
 			key := hashData([]byte(fmt.Sprintf("acc-%d", i)))
-			rawdb.WriteAccountSnapshot(diskdb, key, val)
+			rawdb.WriteAccountSnapshot(helper.diskdb, key, val)
 		}
 	}
-	root, _, _ := accTrie.Commit(nil)
-	t.Logf("root: %x", root)
-	triedb.Commit(root, false, nil)
-
-	snap := generateSnapshot(diskdb, triedb, 16, root)
+	root, snap := helper.CommitAndGenerate()
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -675,31 +620,22 @@ func TestGenerateWithExtraBeforeAndAfter(t *testing.T) {
 	if false {
 		enableLogging()
 	}
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	accTrie, _ := trie.New(common.Hash{}, triedb)
+	helper := newHelper()
 	{
 		acc := &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
 		val, _ := rlp.EncodeToBytes(acc)
-		accTrie.Update(common.HexToHash("0x03").Bytes(), val)
-		accTrie.Update(common.HexToHash("0x07").Bytes(), val)
+		helper.accTrie.Update(common.HexToHash("0x03").Bytes(), val)
+		helper.accTrie.Update(common.HexToHash("0x07").Bytes(), val)
 
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x01"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x02"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x03"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x04"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x05"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x06"), val)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x07"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x01"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x02"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x03"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x04"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x05"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x06"), val)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x07"), val)
 	}
-
-	root, _, _ := accTrie.Commit(nil)
-	t.Logf("root: %x", root)
-	triedb.Commit(root, false, nil)
-
-	snap := generateSnapshot(diskdb, triedb, 16, root)
+	root, snap := helper.CommitAndGenerate()
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -721,29 +657,20 @@ func TestGenerateWithMalformedSnapdata(t *testing.T) {
 	if false {
 		enableLogging()
 	}
-	var (
-		diskdb = memorydb.New()
-		triedb = trie.NewDatabase(diskdb)
-	)
-	accTrie, _ := trie.New(common.Hash{}, triedb)
+	helper := newHelper()
 	{
 		acc := &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()}
 		val, _ := rlp.EncodeToBytes(acc)
-		accTrie.Update(common.HexToHash("0x03").Bytes(), val)
+		helper.accTrie.Update(common.HexToHash("0x03").Bytes(), val)
 
 		junk := make([]byte, 100)
 		copy(junk, []byte{0xde, 0xad})
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x02"), junk)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x03"), junk)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x04"), junk)
-		rawdb.WriteAccountSnapshot(diskdb, common.HexToHash("0x05"), junk)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x02"), junk)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x03"), junk)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x04"), junk)
+		rawdb.WriteAccountSnapshot(helper.diskdb, common.HexToHash("0x05"), junk)
 	}
-
-	root, _, _ := accTrie.Commit(nil)
-	t.Logf("root: %x", root)
-	triedb.Commit(root, false, nil)
-
-	snap := generateSnapshot(diskdb, triedb, 16, root)
+	root, snap := helper.CommitAndGenerate()
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -757,7 +684,7 @@ func TestGenerateWithMalformedSnapdata(t *testing.T) {
 	snap.genAbort <- stop
 	<-stop
 	// If we now inspect the snap db, there should exist no extraneous storage items
-	if data := rawdb.ReadStorageSnapshot(diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
+	if data := rawdb.ReadStorageSnapshot(helper.diskdb, hashData([]byte("acc-2")), hashData([]byte("b-key-1"))); data != nil {
 		t.Fatalf("expected slot to be removed, got %v", string(data))
 	}
 }
@@ -767,13 +694,13 @@ func TestGenerateFromEmptySnap(t *testing.T) {
 	accountCheckRange = 10
 	storageCheckRange = 20
 	helper := newHelper()
-	stRoot := helper.makeStorageTrie([]string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
 	// Add 1K accounts to the trie
 	for i := 0; i < 400; i++ {
+		stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte(fmt.Sprintf("acc-%d", i))), []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 		helper.addTrieAccount(fmt.Sprintf("acc-%d", i),
 			&Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 	}
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	t.Logf("Root: %#x\n", root) // Root: 0x6f7af6d2e1a1bf2b84a3beb3f8b64388465fbc1e274ca5d5d3fc787ca78f59e4
 
 	select {
@@ -802,12 +729,12 @@ func TestGenerateWithIncompleteStorage(t *testing.T) {
 	helper := newHelper()
 	stKeys := []string{"1", "2", "3", "4", "5", "6", "7", "8"}
 	stVals := []string{"v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8"}
-	stRoot := helper.makeStorageTrie(stKeys, stVals)
 	// We add 8 accounts, each one is missing exactly one of the storage slots. This means
 	// we don't have to order the keys and figure out exactly which hash-key winds up
 	// on the sensitive spots at the boundaries
 	for i := 0; i < 8; i++ {
 		accKey := fmt.Sprintf("acc-%d", i)
+		stRoot := helper.makeStorageTrie(common.Hash{}, hashData([]byte(accKey)), stKeys, stVals, true)
 		helper.addAccount(accKey, &Account{Balance: big.NewInt(int64(i)), Root: stRoot, CodeHash: emptyCode.Bytes()})
 		var moddedKeys []string
 		var moddedVals []string
@@ -819,8 +746,7 @@ func TestGenerateWithIncompleteStorage(t *testing.T) {
 		}
 		helper.addSnapStorage(accKey, moddedKeys, moddedVals)
 	}
-
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	t.Logf("Root: %#x\n", root) // Root: 0xca73f6f05ba4ca3024ef340ef3dfca8fdabc1b677ff13f5a9571fd49c16e67ff
 
 	select {
@@ -899,7 +825,7 @@ func populateDangling(disk ethdb.KeyValueStore) {
 // This test will populate some dangling storages to see if they can be cleaned up.
 func TestGenerateCompleteSnapshotWithDanglingStorage(t *testing.T) {
 	var helper = newHelper()
-	stRoot := helper.makeStorageTrie([]string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
+	stRoot := helper.makeStorageTrie(common.Hash{}, common.Hash{}, []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
 	helper.addAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 	helper.addAccount("acc-2", &Account{Balance: big.NewInt(1), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
@@ -910,7 +836,7 @@ func TestGenerateCompleteSnapshotWithDanglingStorage(t *testing.T) {
 
 	populateDangling(helper.diskdb)
 
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded
@@ -932,7 +858,7 @@ func TestGenerateCompleteSnapshotWithDanglingStorage(t *testing.T) {
 // This test will populate some dangling storages to see if they can be cleaned up.
 func TestGenerateBrokenSnapshotWithDanglingStorage(t *testing.T) {
 	var helper = newHelper()
-	stRoot := helper.makeStorageTrie([]string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"})
+	stRoot := helper.makeStorageTrie(common.Hash{}, common.Hash{}, []string{"key-1", "key-2", "key-3"}, []string{"val-1", "val-2", "val-3"}, true)
 
 	helper.addTrieAccount("acc-1", &Account{Balance: big.NewInt(1), Root: stRoot, CodeHash: emptyCode.Bytes()})
 	helper.addTrieAccount("acc-2", &Account{Balance: big.NewInt(2), Root: emptyRoot.Bytes(), CodeHash: emptyCode.Bytes()})
@@ -940,7 +866,7 @@ func TestGenerateBrokenSnapshotWithDanglingStorage(t *testing.T) {
 
 	populateDangling(helper.diskdb)
 
-	root, snap := helper.Generate()
+	root, snap := helper.CommitAndGenerate()
 	select {
 	case <-snap.genPending:
 		// Snapshot generation succeeded

core/state/state_object.go

@@ -154,7 +154,7 @@ func (s *stateObject) getTrie(db Database) Trie {
 		if s.data.Root != emptyRoot && s.db.prefetcher != nil {
 			// When the miner is creating the pending state, there is no
 			// prefetcher
-			s.trie = s.db.prefetcher.trie(s.data.Root)
+			s.trie = s.db.prefetcher.trie(s.addrHash, s.data.Root)
 		}
 		if s.trie == nil {
 			var err error
@@ -295,7 +295,7 @@ func (s *stateObject) finalise(prefetch bool) {
 		}
 	}
 	if s.db.prefetcher != nil && prefetch && len(slotsToPrefetch) > 0 && s.data.Root != emptyRoot {
-		s.db.prefetcher.prefetch(s.data.Root, slotsToPrefetch)
+		s.db.prefetcher.prefetch(s.addrHash, s.data.Root, slotsToPrefetch)
 	}
 	if len(s.dirtyStorage) > 0 {
 		s.dirtyStorage = make(Storage)
@@ -352,7 +352,7 @@ func (s *stateObject) updateTrie(db Database) Trie {
 		usedStorage = append(usedStorage, common.CopyBytes(key[:])) // Copy needed for closure
 	}
 	if s.db.prefetcher != nil {
-		s.db.prefetcher.used(s.data.Root, usedStorage)
+		s.db.prefetcher.used(s.addrHash, s.data.Root, usedStorage)
 	}
 	if len(s.pendingStorage) > 0 {
 		s.pendingStorage = make(Storage)

core/state/statedb.go

@@ -62,11 +62,14 @@ func (n *proofList) Delete(key []byte) error {
 // * Contracts
 // * Accounts
 type StateDB struct {
-	db           Database
-	prefetcher   *triePrefetcher
-	originalRoot common.Hash // The pre-state root, before any changes were made
-	trie         Trie
-	hasher       crypto.KeccakState
+	db         Database
+	prefetcher *triePrefetcher
+	trie       Trie
+	hasher     crypto.KeccakState
+
+	// originalRoot is the pre-state root, before any changes were made.
+	// It will be updated when the Commit is called.
+	originalRoot common.Hash
 
 	snaps         *snapshot.Tree
 	snap          snapshot.Snapshot
@@ -648,6 +651,7 @@ func (s *StateDB) Copy() *StateDB {
 	state := &StateDB{
 		db:                  s.db,
 		trie:                s.db.CopyTrie(s.trie),
+		originalRoot:        s.originalRoot,
 		stateObjects:        make(map[common.Address]*stateObject, len(s.journal.dirties)),
 		stateObjectsPending: make(map[common.Address]struct{}, len(s.stateObjectsPending)),
 		stateObjectsDirty:   make(map[common.Address]struct{}, len(s.journal.dirties)),
@@ -810,7 +814,7 @@ func (s *StateDB) Finalise(deleteEmptyObjects bool) {
 		addressesToPrefetch = append(addressesToPrefetch, common.CopyBytes(addr[:])) // Copy needed for closure
 	}
 	if s.prefetcher != nil && len(addressesToPrefetch) > 0 {
-		s.prefetcher.prefetch(s.originalRoot, addressesToPrefetch)
+		s.prefetcher.prefetch(common.Hash{}, s.originalRoot, addressesToPrefetch)
 	}
 	// Invalidate journal because reverting across transactions is not allowed.
 	s.clearJournalAndRefund()
@@ -851,7 +855,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 	// _untouched_. We can check with the prefetcher, if it can give us a trie
 	// which has the same root, but also has some content loaded into it.
 	if prefetcher != nil {
-		if trie := prefetcher.trie(s.originalRoot); trie != nil {
+		if trie := prefetcher.trie(common.Hash{}, s.originalRoot); trie != nil {
 			s.trie = trie
 		}
 	}
@@ -867,7 +871,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 		usedAddrs = append(usedAddrs, common.CopyBytes(addr[:])) // Copy needed for closure
 	}
 	if prefetcher != nil {
-		prefetcher.used(s.originalRoot, usedAddrs)
+		prefetcher.used(common.Hash{}, s.originalRoot, usedAddrs)
 	}
 	if len(s.stateObjectsPending) > 0 {
 		s.stateObjectsPending = make(map[common.Address]struct{})
@@ -980,6 +984,7 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		}
 		s.snap, s.snapDestructs, s.snapAccounts, s.snapStorage = nil, nil, nil, nil
 	}
+	s.originalRoot = root
 	return root, err
 }
 

core/state/sync_test.go

@@ -104,7 +104,7 @@ func checkTrieConsistency(db ethdb.Database, root common.Hash) error {
 	if v, _ := db.Get(root[:]); v == nil {
 		return nil // Consider a non existent state consistent.
 	}
-	trie, err := trie.New(root, trie.NewDatabase(db))
+	trie, err := trie.New(common.Hash{}, root, trie.NewDatabase(db))
 	if err != nil {
 		return err
 	}
@@ -166,7 +166,7 @@ func testIterativeStateSync(t *testing.T, count int, commit bool, bypath bool) {
 	if commit {
 		srcDb.TrieDB().Commit(srcRoot, false, nil)
 	}
-	srcTrie, _ := trie.New(srcRoot, srcDb.TrieDB())
+	srcTrie, _ := trie.New(common.Hash{}, srcRoot, srcDb.TrieDB())
 
 	// Create a destination state and sync with the scheduler
 	dstDb := rawdb.NewMemoryDatabase()
@@ -207,7 +207,7 @@ func testIterativeStateSync(t *testing.T, count int, commit bool, bypath bool) {
 				if err := rlp.DecodeBytes(srcTrie.Get(path[0]), &acc); err != nil {
 					t.Fatalf("failed to decode account on path %x: %v", path, err)
 				}
-				stTrie, err := trie.New(acc.Root, srcDb.TrieDB())
+				stTrie, err := trie.New(common.BytesToHash(path[0]), acc.Root, srcDb.TrieDB())
 				if err != nil {
 					t.Fatalf("failed to retriev storage trie for path %x: %v", path, err)
 				}

core/state/trie_prefetcher.go

@@ -25,7 +25,7 @@ import (
 )
 
 var (
-	// triePrefetchMetricsPrefix is the prefix under which to publis the metrics.
+	// triePrefetchMetricsPrefix is the prefix under which to publish the metrics.
 	triePrefetchMetricsPrefix = "trie/prefetch/"
 )
 
@@ -35,10 +35,10 @@ var (
 //
 // Note, the prefetcher's API is not thread safe.
 type triePrefetcher struct {
-	db       Database                    // Database to fetch trie nodes through
-	root     common.Hash                 // Root hash of theaccount trie for metrics
-	fetches  map[common.Hash]Trie        // Partially or fully fetcher tries
-	fetchers map[common.Hash]*subfetcher // Subfetchers for each trie
+	db       Database               // Database to fetch trie nodes through
+	root     common.Hash            // Root hash of the account trie for metrics
+	fetches  map[string]Trie        // Partially or fully fetcher tries
+	fetchers map[string]*subfetcher // Subfetchers for each trie
 
 	deliveryMissMeter metrics.Meter
 	accountLoadMeter  metrics.Meter
@@ -51,13 +51,12 @@ type triePrefetcher struct {
 	storageWasteMeter metrics.Meter
 }
 
-// newTriePrefetcher
 func newTriePrefetcher(db Database, root common.Hash, namespace string) *triePrefetcher {
 	prefix := triePrefetchMetricsPrefix + namespace
 	p := &triePrefetcher{
 		db:       db,
 		root:     root,
-		fetchers: make(map[common.Hash]*subfetcher), // Active prefetchers use the fetchers map
+		fetchers: make(map[string]*subfetcher), // Active prefetchers use the fetchers map
 
 		deliveryMissMeter: metrics.GetOrRegisterMeter(prefix+"/deliverymiss", nil),
 		accountLoadMeter:  metrics.GetOrRegisterMeter(prefix+"/account/load", nil),
@@ -112,7 +111,7 @@ func (p *triePrefetcher) copy() *triePrefetcher {
 	copy := &triePrefetcher{
 		db:      p.db,
 		root:    p.root,
-		fetches: make(map[common.Hash]Trie), // Active prefetchers use the fetches map
+		fetches: make(map[string]Trie), // Active prefetchers use the fetches map
 
 		deliveryMissMeter: p.deliveryMissMeter,
 		accountLoadMeter:  p.accountLoadMeter,
@@ -132,33 +131,35 @@ func (p *triePrefetcher) copy() *triePrefetcher {
 		return copy
 	}
 	// Otherwise we're copying an active fetcher, retrieve the current states
-	for root, fetcher := range p.fetchers {
-		copy.fetches[root] = fetcher.peek()
+	for id, fetcher := range p.fetchers {
+		copy.fetches[id] = fetcher.peek()
 	}
 	return copy
 }
 
 // prefetch schedules a batch of trie items to prefetch.
-func (p *triePrefetcher) prefetch(root common.Hash, keys [][]byte) {
+func (p *triePrefetcher) prefetch(owner common.Hash, root common.Hash, keys [][]byte) {
 	// If the prefetcher is an inactive one, bail out
 	if p.fetches != nil {
 		return
 	}
 	// Active fetcher, schedule the retrievals
-	fetcher := p.fetchers[root]
+	id := p.trieID(owner, root)
+	fetcher := p.fetchers[id]
 	if fetcher == nil {
-		fetcher = newSubfetcher(p.db, root)
-		p.fetchers[root] = fetcher
+		fetcher = newSubfetcher(p.db, owner, root)
+		p.fetchers[id] = fetcher
 	}
 	fetcher.schedule(keys)
 }
 
 // trie returns the trie matching the root hash, or nil if the prefetcher doesn't
 // have it.
-func (p *triePrefetcher) trie(root common.Hash) Trie {
+func (p *triePrefetcher) trie(owner common.Hash, root common.Hash) Trie {
 	// If the prefetcher is inactive, return from existing deep copies
+	id := p.trieID(owner, root)
 	if p.fetches != nil {
-		trie := p.fetches[root]
+		trie := p.fetches[id]
 		if trie == nil {
 			p.deliveryMissMeter.Mark(1)
 			return nil
@@ -166,7 +167,7 @@ func (p *triePrefetcher) trie(root common.Hash) Trie {
 		return p.db.CopyTrie(trie)
 	}
 	// Otherwise the prefetcher is active, bail if no trie was prefetched for this root
-	fetcher := p.fetchers[root]
+	fetcher := p.fetchers[id]
 	if fetcher == nil {
 		p.deliveryMissMeter.Mark(1)
 		return nil
@@ -185,20 +186,26 @@ func (p *triePrefetcher) trie(root common.Hash) Trie {
 
 // used marks a batch of state items used to allow creating statistics as to
 // how useful or wasteful the prefetcher is.
-func (p *triePrefetcher) used(root common.Hash, used [][]byte) {
-	if fetcher := p.fetchers[root]; fetcher != nil {
+func (p *triePrefetcher) used(owner common.Hash, root common.Hash, used [][]byte) {
+	if fetcher := p.fetchers[p.trieID(owner, root)]; fetcher != nil {
 		fetcher.used = used
 	}
 }
 
+// trieID returns an unique trie identifier consists the trie owner and root hash.
+func (p *triePrefetcher) trieID(owner common.Hash, root common.Hash) string {
+	return string(append(owner.Bytes(), root.Bytes()...))
+}
+
 // subfetcher is a trie fetcher goroutine responsible for pulling entries for a
 // single trie. It is spawned when a new root is encountered and lives until the
 // main prefetcher is paused and either all requested items are processed or if
 // the trie being worked on is retrieved from the prefetcher.
 type subfetcher struct {
-	db   Database    // Database to load trie nodes through
-	root common.Hash // Root hash of the trie to prefetch
-	trie Trie        // Trie being populated with nodes
+	db    Database    // Database to load trie nodes through
+	owner common.Hash // Owner of the trie, usually account hash
+	root  common.Hash // Root hash of the trie to prefetch
+	trie  Trie        // Trie being populated with nodes
 
 	tasks [][]byte   // Items queued up for retrieval
 	lock  sync.Mutex // Lock protecting the task queue
@@ -215,15 +222,16 @@ type subfetcher struct {
 
 // newSubfetcher creates a goroutine to prefetch state items belonging to a
 // particular root hash.
-func newSubfetcher(db Database, root common.Hash) *subfetcher {
+func newSubfetcher(db Database, owner common.Hash, root common.Hash) *subfetcher {
 	sf := &subfetcher{
-		db:   db,
-		root: root,
-		wake: make(chan struct{}, 1),
-		stop: make(chan struct{}),
-		term: make(chan struct{}),
-		copy: make(chan chan Trie),
-		seen: make(map[string]struct{}),
+		db:    db,
+		owner: owner,
+		root:  root,
+		wake:  make(chan struct{}, 1),
+		stop:  make(chan struct{}),
+		term:  make(chan struct{}),
+		copy:  make(chan chan Trie),
+		seen:  make(map[string]struct{}),
 	}
 	go sf.loop()
 	return sf
@@ -279,13 +287,21 @@ func (sf *subfetcher) loop() {
 	defer close(sf.term)
 
 	// Start by opening the trie and stop processing if it fails
-	trie, err := sf.db.OpenTrie(sf.root)
-	if err != nil {
-		log.Warn("Trie prefetcher failed opening trie", "root", sf.root, "err", err)
-		return
+	if sf.owner == (common.Hash{}) {
+		trie, err := sf.db.OpenTrie(sf.root)
+		if err != nil {
+			log.Warn("Trie prefetcher failed opening trie", "root", sf.root, "err", err)
+			return
+		}
+		sf.trie = trie
+	} else {
+		trie, err := sf.db.OpenStorageTrie(sf.owner, sf.root)
+		if err != nil {
+			log.Warn("Trie prefetcher failed opening trie", "root", sf.root, "err", err)
+			return
+		}
+		sf.trie = trie
 	}
-	sf.trie = trie
-
 	// Trie opened successfully, keep prefetching items
 	for {
 		select {

core/state/trie_prefetcher_test.go

@@ -47,20 +47,20 @@ func TestCopyAndClose(t *testing.T) {
 	db := filledStateDB()
 	prefetcher := newTriePrefetcher(db.db, db.originalRoot, "")
 	skey := common.HexToHash("aaa")
-	prefetcher.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	prefetcher.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
+	prefetcher.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	prefetcher.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
 	time.Sleep(1 * time.Second)
-	a := prefetcher.trie(db.originalRoot)
-	prefetcher.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	b := prefetcher.trie(db.originalRoot)
+	a := prefetcher.trie(common.Hash{}, db.originalRoot)
+	prefetcher.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	b := prefetcher.trie(common.Hash{}, db.originalRoot)
 	cpy := prefetcher.copy()
-	cpy.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	cpy.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	c := cpy.trie(db.originalRoot)
+	cpy.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	cpy.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	c := cpy.trie(common.Hash{}, db.originalRoot)
 	prefetcher.close()
 	cpy2 := cpy.copy()
-	cpy2.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	d := cpy2.trie(db.originalRoot)
+	cpy2.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	d := cpy2.trie(common.Hash{}, db.originalRoot)
 	cpy.close()
 	cpy2.close()
 	if a.Hash() != b.Hash() || a.Hash() != c.Hash() || a.Hash() != d.Hash() {
@@ -72,10 +72,10 @@ func TestUseAfterClose(t *testing.T) {
 	db := filledStateDB()
 	prefetcher := newTriePrefetcher(db.db, db.originalRoot, "")
 	skey := common.HexToHash("aaa")
-	prefetcher.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
-	a := prefetcher.trie(db.originalRoot)
+	prefetcher.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
+	a := prefetcher.trie(common.Hash{}, db.originalRoot)
 	prefetcher.close()
-	b := prefetcher.trie(db.originalRoot)
+	b := prefetcher.trie(common.Hash{}, db.originalRoot)
 	if a == nil {
 		t.Fatal("Prefetching before close should not return nil")
 	}
@@ -88,13 +88,13 @@ func TestCopyClose(t *testing.T) {
 	db := filledStateDB()
 	prefetcher := newTriePrefetcher(db.db, db.originalRoot, "")
 	skey := common.HexToHash("aaa")
-	prefetcher.prefetch(db.originalRoot, [][]byte{skey.Bytes()})
+	prefetcher.prefetch(common.Hash{}, db.originalRoot, [][]byte{skey.Bytes()})
 	cpy := prefetcher.copy()
-	a := prefetcher.trie(db.originalRoot)
-	b := cpy.trie(db.originalRoot)
+	a := prefetcher.trie(common.Hash{}, db.originalRoot)
+	b := cpy.trie(common.Hash{}, db.originalRoot)
 	prefetcher.close()
-	c := prefetcher.trie(db.originalRoot)
-	d := cpy.trie(db.originalRoot)
+	c := prefetcher.trie(common.Hash{}, db.originalRoot)
+	d := cpy.trie(common.Hash{}, db.originalRoot)
 	if a == nil {
 		t.Fatal("Prefetching before close should not return nil")
 	}

core/types/hashing_test.go

@@ -39,8 +39,7 @@ func TestDeriveSha(t *testing.T) {
 		t.Fatal(err)
 	}
 	for len(txs) < 1000 {
-		tr, _ := trie.New(common.Hash{}, trie.NewDatabase(rawdb.NewMemoryDatabase()))
-		exp := types.DeriveSha(txs, tr)
+		exp := types.DeriveSha(txs, trie.NewEmpty(trie.NewDatabase(rawdb.NewMemoryDatabase())))
 		got := types.DeriveSha(txs, trie.NewStackTrie(nil))
 		if !bytes.Equal(got[:], exp[:]) {
 			t.Fatalf("%d txs: got %x exp %x", len(txs), got, exp)
@@ -87,8 +86,7 @@ func BenchmarkDeriveSha200(b *testing.B) {
 		b.ResetTimer()
 		b.ReportAllocs()
 		for i := 0; i < b.N; i++ {
-			tr, _ := trie.New(common.Hash{}, trie.NewDatabase(rawdb.NewMemoryDatabase()))
-			exp = types.DeriveSha(txs, tr)
+			exp = types.DeriveSha(txs, trie.NewEmpty(trie.NewDatabase(rawdb.NewMemoryDatabase())))
 		}
 	})
 
@@ -109,8 +107,7 @@ func TestFuzzDeriveSha(t *testing.T) {
 	rndSeed := mrand.Int()
 	for i := 0; i < 10; i++ {
 		seed := rndSeed + i
-		tr, _ := trie.New(common.Hash{}, trie.NewDatabase(rawdb.NewMemoryDatabase()))
-		exp := types.DeriveSha(newDummy(i), tr)
+		exp := types.DeriveSha(newDummy(i), trie.NewEmpty(trie.NewDatabase(rawdb.NewMemoryDatabase())))
 		got := types.DeriveSha(newDummy(i), trie.NewStackTrie(nil))
 		if !bytes.Equal(got[:], exp[:]) {
 			printList(newDummy(seed))
@@ -138,8 +135,7 @@ func TestDerivableList(t *testing.T) {
 		},
 	}
 	for i, tc := range tcs[1:] {
-		tr, _ := trie.New(common.Hash{}, trie.NewDatabase(rawdb.NewMemoryDatabase()))
-		exp := types.DeriveSha(flatList(tc), tr)
+		exp := types.DeriveSha(flatList(tc), trie.NewEmpty(trie.NewDatabase(rawdb.NewMemoryDatabase())))
 		got := types.DeriveSha(flatList(tc), trie.NewStackTrie(nil))
 		if !bytes.Equal(got[:], exp[:]) {
 			t.Fatalf("case %d: got %x exp %x", i, got, exp)