iterator.go 4.86 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
// Copyright 2015 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package state

import (
	"bytes"
	"fmt"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
)

// NodeIterator is an iterator to traverse the entire state trie post-order,
// including all of the contract code and contract state tries.
type NodeIterator struct {
	state *StateDB // State being iterated

33 34
	stateIt trie.NodeIterator // Primary iterator for the global state trie
	dataIt  trie.NodeIterator // Secondary iterator for the data trie of a contract
35

36 37 38
	accountHash common.Hash // Hash of the node containing the account
	codeHash    common.Hash // Hash of the contract source code
	code        []byte      // Source code associated with a contract
39

40 41
	Hash   common.Hash // Hash of the current entry being iterated (nil if not standalone)
	Parent common.Hash // Hash of the first full ancestor node (nil if current is the root)
42 43

	Error error // Failure set in case of an internal error in the iterator
44 45 46 47 48 49 50 51 52 53
}

// NewNodeIterator creates an post-order state node iterator.
func NewNodeIterator(state *StateDB) *NodeIterator {
	return &NodeIterator{
		state: state,
	}
}

// Next moves the iterator to the next node, returning whether there are any
54 55
// further nodes. In case of an internal error this method returns false and
// sets the Error field to the encountered failure.
56
func (it *NodeIterator) Next() bool {
57 58 59 60 61 62 63 64 65
	// If the iterator failed previously, don't do anything
	if it.Error != nil {
		return false
	}
	// Otherwise step forward with the iterator and report any errors
	if err := it.step(); err != nil {
		it.Error = err
		return false
	}
66 67 68 69
	return it.retrieve()
}

// step moves the iterator to the next entry of the state trie.
70
func (it *NodeIterator) step() error {
71 72
	// Abort if we reached the end of the iteration
	if it.state == nil {
73
		return nil
74 75 76
	}
	// Initialize the iterator if we've just started
	if it.stateIt == nil {
77
		it.stateIt = it.state.trie.NodeIterator(nil)
78 79 80
	}
	// If we had data nodes previously, we surely have at least state nodes
	if it.dataIt != nil {
81 82 83
		if cont := it.dataIt.Next(true); !cont {
			if it.dataIt.Error() != nil {
				return it.dataIt.Error()
84
			}
85 86
			it.dataIt = nil
		}
87
		return nil
88 89 90 91
	}
	// If we had source code previously, discard that
	if it.code != nil {
		it.code = nil
92
		return nil
93 94
	}
	// Step to the next state trie node, terminating if we're out of nodes
95 96 97
	if cont := it.stateIt.Next(true); !cont {
		if it.stateIt.Error() != nil {
			return it.stateIt.Error()
98
		}
99
		it.state, it.stateIt = nil, nil
100
		return nil
101 102
	}
	// If the state trie node is an internal entry, leave as is
103
	if !it.stateIt.Leaf() {
104
		return nil
105 106
	}
	// Otherwise we've reached an account node, initiate data iteration
107
	var account Account
108
	if err := rlp.Decode(bytes.NewReader(it.stateIt.LeafBlob()), &account); err != nil {
109
		return err
110
	}
111
	dataTrie, err := it.state.db.OpenStorageTrie(common.BytesToHash(it.stateIt.LeafKey()), account.Root)
112
	if err != nil {
113
		return err
114
	}
115
	it.dataIt = dataTrie.NodeIterator(nil)
116
	if !it.dataIt.Next(true) {
117 118
		it.dataIt = nil
	}
119
	if !bytes.Equal(account.CodeHash, emptyCodeHash) {
120
		it.codeHash = common.BytesToHash(account.CodeHash)
121 122
		addrHash := common.BytesToHash(it.stateIt.LeafKey())
		it.code, err = it.state.db.ContractCode(addrHash, common.BytesToHash(account.CodeHash))
123
		if err != nil {
124
			return fmt.Errorf("code %x: %v", account.CodeHash, err)
125 126
		}
	}
127
	it.accountHash = it.stateIt.Parent()
128
	return nil
129 130 131 132 133 134
}

// retrieve pulls and caches the current state entry the iterator is traversing.
// The method returns whether there are any more data left for inspection.
func (it *NodeIterator) retrieve() bool {
	// Clear out any previously set values
135
	it.Hash = common.Hash{}
136 137 138 139 140 141 142 143

	// If the iteration's done, return no available data
	if it.state == nil {
		return false
	}
	// Otherwise retrieve the current entry
	switch {
	case it.dataIt != nil:
144
		it.Hash, it.Parent = it.dataIt.Hash(), it.dataIt.Parent()
145 146 147
		if it.Parent == (common.Hash{}) {
			it.Parent = it.accountHash
		}
148
	case it.code != nil:
149
		it.Hash, it.Parent = it.codeHash, it.accountHash
150
	case it.stateIt != nil:
151
		it.Hash, it.Parent = it.stateIt.Hash(), it.stateIt.Parent()
152 153 154
	}
	return true
}