github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/core/state/iterator.go (about)

     1  package state
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  
     7  	"github.com/quickchainproject/quickchain/common"
     8  	"github.com/quickchainproject/quickchain/rlp"
     9  	"github.com/quickchainproject/quickchain/trie"
    10  )
    11  
    12  // NodeIterator is an iterator to traverse the entire state trie post-order,
    13  // including all of the contract code and contract state tries.
    14  type NodeIterator struct {
    15  	state *StateDB // State being iterated
    16  
    17  	stateIt trie.NodeIterator // Primary iterator for the global state trie
    18  	dataIt  trie.NodeIterator // Secondary iterator for the data trie of a contract
    19  
    20  	accountHash common.Hash // Hash of the node containing the account
    21  	codeHash    common.Hash // Hash of the contract source code
    22  	code        []byte      // Source code associated with a contract
    23  
    24  	Hash   common.Hash // Hash of the current entry being iterated (nil if not standalone)
    25  	Parent common.Hash // Hash of the first full ancestor node (nil if current is the root)
    26  
    27  	Error error // Failure set in case of an internal error in the iterator
    28  }
    29  
    30  // NewNodeIterator creates an post-order state node iterator.
    31  func NewNodeIterator(state *StateDB) *NodeIterator {
    32  	return &NodeIterator{
    33  		state: state,
    34  	}
    35  }
    36  
    37  // Next moves the iterator to the next node, returning whether there are any
    38  // further nodes. In case of an internal error this method returns false and
    39  // sets the Error field to the encountered failure.
    40  func (it *NodeIterator) Next() bool {
    41  	// If the iterator failed previously, don't do anything
    42  	if it.Error != nil {
    43  		return false
    44  	}
    45  	// Otherwise step forward with the iterator and report any errors
    46  	if err := it.step(); err != nil {
    47  		it.Error = err
    48  		return false
    49  	}
    50  	return it.retrieve()
    51  }
    52  
    53  // step moves the iterator to the next entry of the state trie.
    54  func (it *NodeIterator) step() error {
    55  	// Abort if we reached the end of the iteration
    56  	if it.state == nil {
    57  		return nil
    58  	}
    59  	// Initialize the iterator if we've just started
    60  	if it.stateIt == nil {
    61  		it.stateIt = it.state.trie.NodeIterator(nil)
    62  	}
    63  	// If we had data nodes previously, we surely have at least state nodes
    64  	if it.dataIt != nil {
    65  		if cont := it.dataIt.Next(true); !cont {
    66  			if it.dataIt.Error() != nil {
    67  				return it.dataIt.Error()
    68  			}
    69  			it.dataIt = nil
    70  		}
    71  		return nil
    72  	}
    73  	// If we had source code previously, discard that
    74  	if it.code != nil {
    75  		it.code = nil
    76  		return nil
    77  	}
    78  	// Step to the next state trie node, terminating if we're out of nodes
    79  	if cont := it.stateIt.Next(true); !cont {
    80  		if it.stateIt.Error() != nil {
    81  			return it.stateIt.Error()
    82  		}
    83  		it.state, it.stateIt = nil, nil
    84  		return nil
    85  	}
    86  	// If the state trie node is an internal entry, leave as is
    87  	if !it.stateIt.Leaf() {
    88  		return nil
    89  	}
    90  	// Otherwise we've reached an account node, initiate data iteration
    91  	var account Account
    92  	if err := rlp.Decode(bytes.NewReader(it.stateIt.LeafBlob()), &account); err != nil {
    93  		return err
    94  	}
    95  	dataTrie, err := it.state.db.OpenStorageTrie(common.BytesToHash(it.stateIt.LeafKey()), account.Root)
    96  	if err != nil {
    97  		return err
    98  	}
    99  	it.dataIt = dataTrie.NodeIterator(nil)
   100  	if !it.dataIt.Next(true) {
   101  		it.dataIt = nil
   102  	}
   103  	if !bytes.Equal(account.CodeHash, emptyCodeHash) {
   104  		it.codeHash = common.BytesToHash(account.CodeHash)
   105  		addrHash := common.BytesToHash(it.stateIt.LeafKey())
   106  		it.code, err = it.state.db.ContractCode(addrHash, common.BytesToHash(account.CodeHash))
   107  		if err != nil {
   108  			return fmt.Errorf("code %x: %v", account.CodeHash, err)
   109  		}
   110  	}
   111  	it.accountHash = it.stateIt.Parent()
   112  	return nil
   113  }
   114  
   115  // retrieve pulls and caches the current state entry the iterator is traversing.
   116  // The method returns whether there are any more data left for inspection.
   117  func (it *NodeIterator) retrieve() bool {
   118  	// Clear out any previously set values
   119  	it.Hash = common.Hash{}
   120  
   121  	// If the iteration's done, return no available data
   122  	if it.state == nil {
   123  		return false
   124  	}
   125  	// Otherwise retrieve the current entry
   126  	switch {
   127  	case it.dataIt != nil:
   128  		it.Hash, it.Parent = it.dataIt.Hash(), it.dataIt.Parent()
   129  		if it.Parent == (common.Hash{}) {
   130  			it.Parent = it.accountHash
   131  		}
   132  	case it.code != nil:
   133  		it.Hash, it.Parent = it.codeHash, it.accountHash
   134  	case it.stateIt != nil:
   135  		it.Hash, it.Parent = it.stateIt.Hash(), it.stateIt.Parent()
   136  	}
   137  	return true
   138  }