open trie when first invoke

Author: ucwong

core/state/database.go

@@ -248,3 +248,13 @@ func (db *CachingDB) Snapshot() *snapshot.Tree {
 func (db *CachingDB) SetSnapshot(snap *snapshot.Tree) {
 	db.snap = snap
 }
+
+// mustCopyTrie returns a deep-copied trie.
+func mustCopyTrie(t Trie) Trie {
+	switch t := t.(type) {
+	case *trie.StateTrie:
+		return t.Copy()
+	default:
+		panic(fmt.Errorf("unknown trie type %T", t))
+	}
+}

core/state/iterator.go

@@ -32,6 +32,8 @@ import (
 type NodeIterator struct {
 	state *StateDB // State being iterated
 
+	tr Trie // Primary account trie for traversal
+
 	stateIt trie.NodeIterator // Primary iterator for the global state trie
 	dataIt  trie.NodeIterator // Secondary iterator for the data trie of a contract
 
@@ -74,9 +76,20 @@ func (it *NodeIterator) step() error {
 	if it.state == nil {
 		return nil
 	}
+	if it.tr == nil {
+		tr, err := it.state.db.OpenTrie(it.state.originalRoot)
+		if err != nil {
+			return err
+		}
+		it.tr = tr
+	}
 	// Initialize the iterator if we've just started
 	if it.stateIt == nil {
-		it.stateIt = it.state.trie.NodeIterator(nil)
+		stateIt := it.tr.NodeIterator(nil)
+		if stateIt == nil {
+			return errors.New("state iter is nil")
+		}
+		it.stateIt = stateIt
 	}
 	// If we had data nodes previously, we surely have at least state nodes
 	if it.dataIt != nil {
@@ -111,15 +124,14 @@ func (it *NodeIterator) step() error {
 		return err
 	}
 	// Lookup the preimage of account hash
-	preimage := it.state.trie.GetKey(it.stateIt.LeafKey())
+	preimage := it.tr.GetKey(it.stateIt.LeafKey())
 	if preimage == nil {
 		return errors.New("account address is not available")
 	}
 	address := common.BytesToAddress(preimage)
 
 	// Traverse the storage slots belong to the account
-	dataTrie, err := it.state.db.OpenStorageTrie(it.state.originalRoot, address, account.Root, it.state.trie)
-	//dataTrie, err := it.state.db.OpenStorageTrie(common.BytesToAddress(it.stateIt.LeafKey()), account.Root)
+	dataTrie, err := it.state.db.OpenStorageTrie(it.state.originalRoot, address, account.Root, it.tr)
 	if err != nil {
 		return err
 	}

core/state/statedb.go

@@ -166,13 +166,8 @@ type StateDB struct {
 
 // Create a new state from a given trie.
 func New(root common.Hash, db Database) (*StateDB, error) {
-	tr, err := db.OpenTrie(root)
-	if err != nil {
-		return nil, err
-	}
 	sdb := &StateDB{
 		db:                   db,
-		trie:                 tr,
 		originalRoot:         root,
 		accounts:             make(map[common.Hash][]byte),
 		storages:             make(map[common.Hash]map[common.Hash][]byte),
@@ -671,6 +666,14 @@ func (s *StateDB) getStateObject(addr common.Address) *stateObject {
 	// If snapshot unavailable or reading from it failed, load from the database
 	if data == nil {
 		start := time.Now()
+		if s.trie == nil {
+			tr, err := s.db.OpenTrie(s.originalRoot)
+			if err != nil {
+				return nil
+			}
+			s.trie = tr
+
+		}
 		enc, err := s.trie.TryGet(addr.Bytes())
 		s.AccountReads += time.Since(start)
 		if err != nil {
@@ -782,8 +785,7 @@ func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common
 func (s *StateDB) Copy() *StateDB {
 	// Copy all the basic fields, initialize the memory ones
 	state := &StateDB{
-		db:   s.db,
-		trie: s.db.CopyTrie(s.trie),
+		db: s.db,
 		//hasher:               crypto.NewKeccakState(),
 		originalRoot:         s.originalRoot,
 		accounts:             copySet(s.accounts),
@@ -808,6 +810,9 @@ func (s *StateDB) Copy() *StateDB {
 		// miner to operate trie-backed only.
 		snap: s.snap,
 	}
+	if s.trie != nil {
+		state.trie = mustCopyTrie(s.trie)
+	}
 	// Deep copy cached state objects.
 	for addr, obj := range s.stateObjects {
 		state.stateObjects[addr] = obj.deepCopy(state)
@@ -913,6 +918,20 @@ func (s *StateDB) Finalise(deleteEmptyObjects bool) {
 func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 	// Finalise all the dirty storage states and write them into the tries
 	s.Finalise(deleteEmptyObjects)
+	// Initialize the trie if it's not constructed yet. If the prefetch
+	// is enabled, the trie constructed below will be replaced by the
+	// prefetched one.
+	//
+	// This operation must be done before state object storage hashing,
+	// as it assumes the main trie is already loaded.
+	if s.trie == nil {
+		tr, err := s.db.OpenTrie(s.originalRoot)
+		if err != nil {
+			s.setError(err)
+			return common.Hash{}
+		}
+		s.trie = tr
+	}
 
 	// If there was a trie prefetcher operating, terminate it async so that the
 	// individual storage tries can be updated as soon as the disk load finishes.