github.com/vay007/go-ethereum@v1.9.7/trie/iterator.go (about)

     1  // Copyright 2014 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package trie
    18  
    19  import (
    20  	"bytes"
    21  	"container/heap"
    22  	"errors"
    23  
    24  	"github.com/ethereum/go-ethereum/common"
    25  	"github.com/ethereum/go-ethereum/rlp"
    26  )
    27  
    28  // Iterator is a key-value trie iterator that traverses a Trie.
    29  type Iterator struct {
    30  	nodeIt NodeIterator
    31  
    32  	Key   []byte // Current data key on which the iterator is positioned on
    33  	Value []byte // Current data value on which the iterator is positioned on
    34  	Err   error
    35  }
    36  
    37  // NewIterator creates a new key-value iterator from a node iterator.
    38  // Note that the value returned by the iterator is raw. If the content is encoded
    39  // (e.g. storage value is RLP-encoded), it's caller's duty to decode it.
    40  func NewIterator(it NodeIterator) *Iterator {
    41  	return &Iterator{
    42  		nodeIt: it,
    43  	}
    44  }
    45  
    46  // Next moves the iterator forward one key-value entry.
    47  func (it *Iterator) Next() bool {
    48  	for it.nodeIt.Next(true) {
    49  		if it.nodeIt.Leaf() {
    50  			it.Key = it.nodeIt.LeafKey()
    51  			it.Value = it.nodeIt.LeafBlob()
    52  			return true
    53  		}
    54  	}
    55  	it.Key = nil
    56  	it.Value = nil
    57  	it.Err = it.nodeIt.Error()
    58  	return false
    59  }
    60  
    61  // Prove generates the Merkle proof for the leaf node the iterator is currently
    62  // positioned on.
    63  func (it *Iterator) Prove() [][]byte {
    64  	return it.nodeIt.LeafProof()
    65  }
    66  
    67  // NodeIterator is an iterator to traverse the trie pre-order.
    68  type NodeIterator interface {
    69  	// Next moves the iterator to the next node. If the parameter is false, any child
    70  	// nodes will be skipped.
    71  	Next(bool) bool
    72  
    73  	// Error returns the error status of the iterator.
    74  	Error() error
    75  
    76  	// Hash returns the hash of the current node.
    77  	Hash() common.Hash
    78  
    79  	// Parent returns the hash of the parent of the current node. The hash may be the one
    80  	// grandparent if the immediate parent is an internal node with no hash.
    81  	Parent() common.Hash
    82  
    83  	// Path returns the hex-encoded path to the current node.
    84  	// Callers must not retain references to the return value after calling Next.
    85  	// For leaf nodes, the last element of the path is the 'terminator symbol' 0x10.
    86  	Path() []byte
    87  
    88  	// Leaf returns true iff the current node is a leaf node.
    89  	Leaf() bool
    90  
    91  	// LeafKey returns the key of the leaf. The method panics if the iterator is not
    92  	// positioned at a leaf. Callers must not retain references to the value after
    93  	// calling Next.
    94  	LeafKey() []byte
    95  
    96  	// LeafBlob returns the content of the leaf. The method panics if the iterator
    97  	// is not positioned at a leaf. Callers must not retain references to the value
    98  	// after calling Next.
    99  	LeafBlob() []byte
   100  
   101  	// LeafProof returns the Merkle proof of the leaf. The method panics if the
   102  	// iterator is not positioned at a leaf. Callers must not retain references
   103  	// to the value after calling Next.
   104  	LeafProof() [][]byte
   105  }
   106  
   107  // nodeIteratorState represents the iteration state at one particular node of the
   108  // trie, which can be resumed at a later invocation.
   109  type nodeIteratorState struct {
   110  	hash    common.Hash // Hash of the node being iterated (nil if not standalone)
   111  	node    node        // Trie node being iterated
   112  	parent  common.Hash // Hash of the first full ancestor node (nil if current is the root)
   113  	index   int         // Child to be processed next
   114  	pathlen int         // Length of the path to this node
   115  }
   116  
   117  type nodeIterator struct {
   118  	trie  *Trie                // Trie being iterated
   119  	stack []*nodeIteratorState // Hierarchy of trie nodes persisting the iteration state
   120  	path  []byte               // Path to the current node
   121  	err   error                // Failure set in case of an internal error in the iterator
   122  }
   123  
   124  // errIteratorEnd is stored in nodeIterator.err when iteration is done.
   125  var errIteratorEnd = errors.New("end of iteration")
   126  
   127  // seekError is stored in nodeIterator.err if the initial seek has failed.
   128  type seekError struct {
   129  	key []byte
   130  	err error
   131  }
   132  
   133  func (e seekError) Error() string {
   134  	return "seek error: " + e.err.Error()
   135  }
   136  
   137  func newNodeIterator(trie *Trie, start []byte) NodeIterator {
   138  	if trie.Hash() == emptyState {
   139  		return new(nodeIterator)
   140  	}
   141  	it := &nodeIterator{trie: trie}
   142  	it.err = it.seek(start)
   143  	return it
   144  }
   145  
   146  func (it *nodeIterator) Hash() common.Hash {
   147  	if len(it.stack) == 0 {
   148  		return common.Hash{}
   149  	}
   150  	return it.stack[len(it.stack)-1].hash
   151  }
   152  
   153  func (it *nodeIterator) Parent() common.Hash {
   154  	if len(it.stack) == 0 {
   155  		return common.Hash{}
   156  	}
   157  	return it.stack[len(it.stack)-1].parent
   158  }
   159  
   160  func (it *nodeIterator) Leaf() bool {
   161  	return hasTerm(it.path)
   162  }
   163  
   164  func (it *nodeIterator) LeafKey() []byte {
   165  	if len(it.stack) > 0 {
   166  		if _, ok := it.stack[len(it.stack)-1].node.(valueNode); ok {
   167  			return hexToKeybytes(it.path)
   168  		}
   169  	}
   170  	panic("not at leaf")
   171  }
   172  
   173  func (it *nodeIterator) LeafBlob() []byte {
   174  	if len(it.stack) > 0 {
   175  		if node, ok := it.stack[len(it.stack)-1].node.(valueNode); ok {
   176  			return []byte(node)
   177  		}
   178  	}
   179  	panic("not at leaf")
   180  }
   181  
   182  func (it *nodeIterator) LeafProof() [][]byte {
   183  	if len(it.stack) > 0 {
   184  		if _, ok := it.stack[len(it.stack)-1].node.(valueNode); ok {
   185  			hasher := newHasher(nil)
   186  			defer returnHasherToPool(hasher)
   187  
   188  			proofs := make([][]byte, 0, len(it.stack))
   189  
   190  			for i, item := range it.stack[:len(it.stack)-1] {
   191  				// Gather nodes that end up as hash nodes (or the root)
   192  				node, _, _ := hasher.hashChildren(item.node, nil)
   193  				hashed, _ := hasher.store(node, nil, false)
   194  				if _, ok := hashed.(hashNode); ok || i == 0 {
   195  					enc, _ := rlp.EncodeToBytes(node)
   196  					proofs = append(proofs, enc)
   197  				}
   198  			}
   199  			return proofs
   200  		}
   201  	}
   202  	panic("not at leaf")
   203  }
   204  
   205  func (it *nodeIterator) Path() []byte {
   206  	return it.path
   207  }
   208  
   209  func (it *nodeIterator) Error() error {
   210  	if it.err == errIteratorEnd {
   211  		return nil
   212  	}
   213  	if seek, ok := it.err.(seekError); ok {
   214  		return seek.err
   215  	}
   216  	return it.err
   217  }
   218  
   219  // Next moves the iterator to the next node, returning whether there are any
   220  // further nodes. In case of an internal error this method returns false and
   221  // sets the Error field to the encountered failure. If `descend` is false,
   222  // skips iterating over any subnodes of the current node.
   223  func (it *nodeIterator) Next(descend bool) bool {
   224  	if it.err == errIteratorEnd {
   225  		return false
   226  	}
   227  	if seek, ok := it.err.(seekError); ok {
   228  		if it.err = it.seek(seek.key); it.err != nil {
   229  			return false
   230  		}
   231  	}
   232  	// Otherwise step forward with the iterator and report any errors.
   233  	state, parentIndex, path, err := it.peek(descend)
   234  	it.err = err
   235  	if it.err != nil {
   236  		return false
   237  	}
   238  	it.push(state, parentIndex, path)
   239  	return true
   240  }
   241  
   242  func (it *nodeIterator) seek(prefix []byte) error {
   243  	// The path we're looking for is the hex encoded key without terminator.
   244  	key := keybytesToHex(prefix)
   245  	key = key[:len(key)-1]
   246  	// Move forward until we're just before the closest match to key.
   247  	for {
   248  		state, parentIndex, path, err := it.peek(bytes.HasPrefix(key, it.path))
   249  		if err == errIteratorEnd {
   250  			return errIteratorEnd
   251  		} else if err != nil {
   252  			return seekError{prefix, err}
   253  		} else if bytes.Compare(path, key) >= 0 {
   254  			return nil
   255  		}
   256  		it.push(state, parentIndex, path)
   257  	}
   258  }
   259  
   260  // peek creates the next state of the iterator.
   261  func (it *nodeIterator) peek(descend bool) (*nodeIteratorState, *int, []byte, error) {
   262  	if len(it.stack) == 0 {
   263  		// Initialize the iterator if we've just started.
   264  		root := it.trie.Hash()
   265  		state := &nodeIteratorState{node: it.trie.root, index: -1}
   266  		if root != emptyRoot {
   267  			state.hash = root
   268  		}
   269  		err := state.resolve(it.trie, nil)
   270  		return state, nil, nil, err
   271  	}
   272  	if !descend {
   273  		// If we're skipping children, pop the current node first
   274  		it.pop()
   275  	}
   276  
   277  	// Continue iteration to the next child
   278  	for len(it.stack) > 0 {
   279  		parent := it.stack[len(it.stack)-1]
   280  		ancestor := parent.hash
   281  		if (ancestor == common.Hash{}) {
   282  			ancestor = parent.parent
   283  		}
   284  		state, path, ok := it.nextChild(parent, ancestor)
   285  		if ok {
   286  			if err := state.resolve(it.trie, path); err != nil {
   287  				return parent, &parent.index, path, err
   288  			}
   289  			return state, &parent.index, path, nil
   290  		}
   291  		// No more child nodes, move back up.
   292  		it.pop()
   293  	}
   294  	return nil, nil, nil, errIteratorEnd
   295  }
   296  
   297  func (st *nodeIteratorState) resolve(tr *Trie, path []byte) error {
   298  	if hash, ok := st.node.(hashNode); ok {
   299  		resolved, err := tr.resolveHash(hash, path)
   300  		if err != nil {
   301  			return err
   302  		}
   303  		st.node = resolved
   304  		st.hash = common.BytesToHash(hash)
   305  	}
   306  	return nil
   307  }
   308  
   309  func (it *nodeIterator) nextChild(parent *nodeIteratorState, ancestor common.Hash) (*nodeIteratorState, []byte, bool) {
   310  	switch node := parent.node.(type) {
   311  	case *fullNode:
   312  		// Full node, move to the first non-nil child.
   313  		for i := parent.index + 1; i < len(node.Children); i++ {
   314  			child := node.Children[i]
   315  			if child != nil {
   316  				hash, _ := child.cache()
   317  				state := &nodeIteratorState{
   318  					hash:    common.BytesToHash(hash),
   319  					node:    child,
   320  					parent:  ancestor,
   321  					index:   -1,
   322  					pathlen: len(it.path),
   323  				}
   324  				path := append(it.path, byte(i))
   325  				parent.index = i - 1
   326  				return state, path, true
   327  			}
   328  		}
   329  	case *shortNode:
   330  		// Short node, return the pointer singleton child
   331  		if parent.index < 0 {
   332  			hash, _ := node.Val.cache()
   333  			state := &nodeIteratorState{
   334  				hash:    common.BytesToHash(hash),
   335  				node:    node.Val,
   336  				parent:  ancestor,
   337  				index:   -1,
   338  				pathlen: len(it.path),
   339  			}
   340  			path := append(it.path, node.Key...)
   341  			return state, path, true
   342  		}
   343  	}
   344  	return parent, it.path, false
   345  }
   346  
   347  func (it *nodeIterator) push(state *nodeIteratorState, parentIndex *int, path []byte) {
   348  	it.path = path
   349  	it.stack = append(it.stack, state)
   350  	if parentIndex != nil {
   351  		*parentIndex++
   352  	}
   353  }
   354  
   355  func (it *nodeIterator) pop() {
   356  	parent := it.stack[len(it.stack)-1]
   357  	it.path = it.path[:parent.pathlen]
   358  	it.stack = it.stack[:len(it.stack)-1]
   359  }
   360  
   361  func compareNodes(a, b NodeIterator) int {
   362  	if cmp := bytes.Compare(a.Path(), b.Path()); cmp != 0 {
   363  		return cmp
   364  	}
   365  	if a.Leaf() && !b.Leaf() {
   366  		return -1
   367  	} else if b.Leaf() && !a.Leaf() {
   368  		return 1
   369  	}
   370  	if cmp := bytes.Compare(a.Hash().Bytes(), b.Hash().Bytes()); cmp != 0 {
   371  		return cmp
   372  	}
   373  	if a.Leaf() && b.Leaf() {
   374  		return bytes.Compare(a.LeafBlob(), b.LeafBlob())
   375  	}
   376  	return 0
   377  }
   378  
   379  type differenceIterator struct {
   380  	a, b  NodeIterator // Nodes returned are those in b - a.
   381  	eof   bool         // Indicates a has run out of elements
   382  	count int          // Number of nodes scanned on either trie
   383  }
   384  
   385  // NewDifferenceIterator constructs a NodeIterator that iterates over elements in b that
   386  // are not in a. Returns the iterator, and a pointer to an integer recording the number
   387  // of nodes seen.
   388  func NewDifferenceIterator(a, b NodeIterator) (NodeIterator, *int) {
   389  	a.Next(true)
   390  	it := &differenceIterator{
   391  		a: a,
   392  		b: b,
   393  	}
   394  	return it, &it.count
   395  }
   396  
   397  func (it *differenceIterator) Hash() common.Hash {
   398  	return it.b.Hash()
   399  }
   400  
   401  func (it *differenceIterator) Parent() common.Hash {
   402  	return it.b.Parent()
   403  }
   404  
   405  func (it *differenceIterator) Leaf() bool {
   406  	return it.b.Leaf()
   407  }
   408  
   409  func (it *differenceIterator) LeafKey() []byte {
   410  	return it.b.LeafKey()
   411  }
   412  
   413  func (it *differenceIterator) LeafBlob() []byte {
   414  	return it.b.LeafBlob()
   415  }
   416  
   417  func (it *differenceIterator) LeafProof() [][]byte {
   418  	return it.b.LeafProof()
   419  }
   420  
   421  func (it *differenceIterator) Path() []byte {
   422  	return it.b.Path()
   423  }
   424  
   425  func (it *differenceIterator) Next(bool) bool {
   426  	// Invariants:
   427  	// - We always advance at least one element in b.
   428  	// - At the start of this function, a's path is lexically greater than b's.
   429  	if !it.b.Next(true) {
   430  		return false
   431  	}
   432  	it.count++
   433  
   434  	if it.eof {
   435  		// a has reached eof, so we just return all elements from b
   436  		return true
   437  	}
   438  
   439  	for {
   440  		switch compareNodes(it.a, it.b) {
   441  		case -1:
   442  			// b jumped past a; advance a
   443  			if !it.a.Next(true) {
   444  				it.eof = true
   445  				return true
   446  			}
   447  			it.count++
   448  		case 1:
   449  			// b is before a
   450  			return true
   451  		case 0:
   452  			// a and b are identical; skip this whole subtree if the nodes have hashes
   453  			hasHash := it.a.Hash() == common.Hash{}
   454  			if !it.b.Next(hasHash) {
   455  				return false
   456  			}
   457  			it.count++
   458  			if !it.a.Next(hasHash) {
   459  				it.eof = true
   460  				return true
   461  			}
   462  			it.count++
   463  		}
   464  	}
   465  }
   466  
   467  func (it *differenceIterator) Error() error {
   468  	if err := it.a.Error(); err != nil {
   469  		return err
   470  	}
   471  	return it.b.Error()
   472  }
   473  
   474  type nodeIteratorHeap []NodeIterator
   475  
   476  func (h nodeIteratorHeap) Len() int            { return len(h) }
   477  func (h nodeIteratorHeap) Less(i, j int) bool  { return compareNodes(h[i], h[j]) < 0 }
   478  func (h nodeIteratorHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
   479  func (h *nodeIteratorHeap) Push(x interface{}) { *h = append(*h, x.(NodeIterator)) }
   480  func (h *nodeIteratorHeap) Pop() interface{} {
   481  	n := len(*h)
   482  	x := (*h)[n-1]
   483  	*h = (*h)[0 : n-1]
   484  	return x
   485  }
   486  
   487  type unionIterator struct {
   488  	items *nodeIteratorHeap // Nodes returned are the union of the ones in these iterators
   489  	count int               // Number of nodes scanned across all tries
   490  }
   491  
   492  // NewUnionIterator constructs a NodeIterator that iterates over elements in the union
   493  // of the provided NodeIterators. Returns the iterator, and a pointer to an integer
   494  // recording the number of nodes visited.
   495  func NewUnionIterator(iters []NodeIterator) (NodeIterator, *int) {
   496  	h := make(nodeIteratorHeap, len(iters))
   497  	copy(h, iters)
   498  	heap.Init(&h)
   499  
   500  	ui := &unionIterator{items: &h}
   501  	return ui, &ui.count
   502  }
   503  
   504  func (it *unionIterator) Hash() common.Hash {
   505  	return (*it.items)[0].Hash()
   506  }
   507  
   508  func (it *unionIterator) Parent() common.Hash {
   509  	return (*it.items)[0].Parent()
   510  }
   511  
   512  func (it *unionIterator) Leaf() bool {
   513  	return (*it.items)[0].Leaf()
   514  }
   515  
   516  func (it *unionIterator) LeafKey() []byte {
   517  	return (*it.items)[0].LeafKey()
   518  }
   519  
   520  func (it *unionIterator) LeafBlob() []byte {
   521  	return (*it.items)[0].LeafBlob()
   522  }
   523  
   524  func (it *unionIterator) LeafProof() [][]byte {
   525  	return (*it.items)[0].LeafProof()
   526  }
   527  
   528  func (it *unionIterator) Path() []byte {
   529  	return (*it.items)[0].Path()
   530  }
   531  
   532  // Next returns the next node in the union of tries being iterated over.
   533  //
   534  // It does this by maintaining a heap of iterators, sorted by the iteration
   535  // order of their next elements, with one entry for each source trie. Each
   536  // time Next() is called, it takes the least element from the heap to return,
   537  // advancing any other iterators that also point to that same element. These
   538  // iterators are called with descend=false, since we know that any nodes under
   539  // these nodes will also be duplicates, found in the currently selected iterator.
   540  // Whenever an iterator is advanced, it is pushed back into the heap if it still
   541  // has elements remaining.
   542  //
   543  // In the case that descend=false - eg, we're asked to ignore all subnodes of the
   544  // current node - we also advance any iterators in the heap that have the current
   545  // path as a prefix.
   546  func (it *unionIterator) Next(descend bool) bool {
   547  	if len(*it.items) == 0 {
   548  		return false
   549  	}
   550  
   551  	// Get the next key from the union
   552  	least := heap.Pop(it.items).(NodeIterator)
   553  
   554  	// Skip over other nodes as long as they're identical, or, if we're not descending, as
   555  	// long as they have the same prefix as the current node.
   556  	for len(*it.items) > 0 && ((!descend && bytes.HasPrefix((*it.items)[0].Path(), least.Path())) || compareNodes(least, (*it.items)[0]) == 0) {
   557  		skipped := heap.Pop(it.items).(NodeIterator)
   558  		// Skip the whole subtree if the nodes have hashes; otherwise just skip this node
   559  		if skipped.Next(skipped.Hash() == common.Hash{}) {
   560  			it.count++
   561  			// If there are more elements, push the iterator back on the heap
   562  			heap.Push(it.items, skipped)
   563  		}
   564  	}
   565  	if least.Next(descend) {
   566  		it.count++
   567  		heap.Push(it.items, least)
   568  	}
   569  	return len(*it.items) > 0
   570  }
   571  
   572  func (it *unionIterator) Error() error {
   573  	for i := 0; i < len(*it.items); i++ {
   574  		if err := (*it.items)[i].Error(); err != nil {
   575  			return err
   576  		}
   577  	}
   578  	return nil
   579  }