github.com/reapchain/go-reapchain@v0.2.15-0.20210609012950-9735c110c705/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  )
    26  
    27  var iteratorEnd = errors.New("end of iteration")
    28  
    29  // Iterator is a key-value trie iterator that traverses a Trie.
    30  type Iterator struct {
    31  	nodeIt NodeIterator
    32  
    33  	Key   []byte // Current data key on which the iterator is positioned on
    34  	Value []byte // Current data value on which the iterator is positioned on
    35  }
    36  
    37  // NewIterator creates a new key-value iterator from a node iterator
    38  func NewIterator(it NodeIterator) *Iterator {
    39  	return &Iterator{
    40  		nodeIt: it,
    41  	}
    42  }
    43  
    44  // Next moves the iterator forward one key-value entry.
    45  func (it *Iterator) Next() bool {
    46  	for it.nodeIt.Next(true) {
    47  		if it.nodeIt.Leaf() {
    48  			it.Key = hexToKeybytes(it.nodeIt.Path())
    49  			it.Value = it.nodeIt.LeafBlob()
    50  			return true
    51  		}
    52  	}
    53  	it.Key = nil
    54  	it.Value = nil
    55  	return false
    56  }
    57  
    58  // NodeIterator is an iterator to traverse the trie pre-order.
    59  type NodeIterator interface {
    60  	// Hash returns the hash of the current node
    61  	Hash() common.Hash
    62  	// Parent returns the hash of the parent of the current node
    63  	Parent() common.Hash
    64  	// Leaf returns true iff the current node is a leaf node.
    65  	Leaf() bool
    66  	// LeafBlob returns the contents of the node, if it is a leaf.
    67  	// Callers must not retain references to the return value after calling Next()
    68  	LeafBlob() []byte
    69  	// Path returns the hex-encoded path to the current node.
    70  	// Callers must not retain references to the return value after calling Next()
    71  	Path() []byte
    72  	// Next moves the iterator to the next node. If the parameter is false, any child
    73  	// nodes will be skipped.
    74  	Next(bool) bool
    75  	// Error returns the error status of the iterator.
    76  	Error() error
    77  }
    78  
    79  // nodeIteratorState represents the iteration state at one particular node of the
    80  // trie, which can be resumed at a later invocation.
    81  type nodeIteratorState struct {
    82  	hash    common.Hash // Hash of the node being iterated (nil if not standalone)
    83  	node    node        // Trie node being iterated
    84  	parent  common.Hash // Hash of the first full ancestor node (nil if current is the root)
    85  	index   int         // Child to be processed next
    86  	pathlen int         // Length of the path to this node
    87  }
    88  
    89  type nodeIterator struct {
    90  	trie  *Trie                // Trie being iterated
    91  	stack []*nodeIteratorState // Hierarchy of trie nodes persisting the iteration state
    92  	err   error                // Failure set in case of an internal error in the iterator
    93  	path  []byte               // Path to the current node
    94  }
    95  
    96  func newNodeIterator(trie *Trie, start []byte) NodeIterator {
    97  	if trie.Hash() == emptyState {
    98  		return new(nodeIterator)
    99  	}
   100  	it := &nodeIterator{trie: trie}
   101  	it.seek(start)
   102  	return it
   103  }
   104  
   105  // Hash returns the hash of the current node
   106  func (it *nodeIterator) Hash() common.Hash {
   107  	if len(it.stack) == 0 {
   108  		return common.Hash{}
   109  	}
   110  
   111  	return it.stack[len(it.stack)-1].hash
   112  }
   113  
   114  // Parent returns the hash of the parent node
   115  func (it *nodeIterator) Parent() common.Hash {
   116  	if len(it.stack) == 0 {
   117  		return common.Hash{}
   118  	}
   119  
   120  	return it.stack[len(it.stack)-1].parent
   121  }
   122  
   123  // Leaf returns true if the current node is a leaf
   124  func (it *nodeIterator) Leaf() bool {
   125  	if len(it.stack) == 0 {
   126  		return false
   127  	}
   128  
   129  	_, ok := it.stack[len(it.stack)-1].node.(valueNode)
   130  	return ok
   131  }
   132  
   133  // LeafBlob returns the data for the current node, if it is a leaf
   134  func (it *nodeIterator) LeafBlob() []byte {
   135  	if len(it.stack) == 0 {
   136  		return nil
   137  	}
   138  
   139  	if node, ok := it.stack[len(it.stack)-1].node.(valueNode); ok {
   140  		return []byte(node)
   141  	}
   142  	return nil
   143  }
   144  
   145  // Path returns the hex-encoded path to the current node
   146  func (it *nodeIterator) Path() []byte {
   147  	return it.path
   148  }
   149  
   150  // Error returns the error set in case of an internal error in the iterator
   151  func (it *nodeIterator) Error() error {
   152  	if it.err == iteratorEnd {
   153  		return nil
   154  	}
   155  	return it.err
   156  }
   157  
   158  // Next moves the iterator to the next node, returning whether there are any
   159  // further nodes. In case of an internal error this method returns false and
   160  // sets the Error field to the encountered failure. If `descend` is false,
   161  // skips iterating over any subnodes of the current node.
   162  func (it *nodeIterator) Next(descend bool) bool {
   163  	if it.err != nil {
   164  		return false
   165  	}
   166  	// Otherwise step forward with the iterator and report any errors
   167  	state, parentIndex, path, err := it.peek(descend)
   168  	if err != nil {
   169  		it.err = err
   170  		return false
   171  	}
   172  	it.push(state, parentIndex, path)
   173  	return true
   174  }
   175  
   176  func (it *nodeIterator) seek(prefix []byte) {
   177  	// The path we're looking for is the hex encoded key without terminator.
   178  	key := keybytesToHex(prefix)
   179  	key = key[:len(key)-1]
   180  	// Move forward until we're just before the closest match to key.
   181  	for {
   182  		state, parentIndex, path, err := it.peek(bytes.HasPrefix(key, it.path))
   183  		if err != nil || bytes.Compare(path, key) >= 0 {
   184  			it.err = err
   185  			return
   186  		}
   187  		it.push(state, parentIndex, path)
   188  	}
   189  }
   190  
   191  // peek creates the next state of the iterator.
   192  func (it *nodeIterator) peek(descend bool) (*nodeIteratorState, *int, []byte, error) {
   193  	if len(it.stack) == 0 {
   194  		// Initialize the iterator if we've just started.
   195  		root := it.trie.Hash()
   196  		state := &nodeIteratorState{node: it.trie.root, index: -1}
   197  		if root != emptyRoot {
   198  			state.hash = root
   199  		}
   200  		return state, nil, nil, nil
   201  	}
   202  	if !descend {
   203  		// If we're skipping children, pop the current node first
   204  		it.pop()
   205  	}
   206  
   207  	// Continue iteration to the next child
   208  	for {
   209  		if len(it.stack) == 0 {
   210  			return nil, nil, nil, iteratorEnd
   211  		}
   212  		parent := it.stack[len(it.stack)-1]
   213  		ancestor := parent.hash
   214  		if (ancestor == common.Hash{}) {
   215  			ancestor = parent.parent
   216  		}
   217  		if node, ok := parent.node.(*fullNode); ok {
   218  			// Full node, move to the first non-nil child.
   219  			for i := parent.index + 1; i < len(node.Children); i++ {
   220  				child := node.Children[i]
   221  				if child != nil {
   222  					hash, _ := child.cache()
   223  					state := &nodeIteratorState{
   224  						hash:    common.BytesToHash(hash),
   225  						node:    child,
   226  						parent:  ancestor,
   227  						index:   -1,
   228  						pathlen: len(it.path),
   229  					}
   230  					path := append(it.path, byte(i))
   231  					parent.index = i - 1
   232  					return state, &parent.index, path, nil
   233  				}
   234  			}
   235  		} else if node, ok := parent.node.(*shortNode); ok {
   236  			// Short node, return the pointer singleton child
   237  			if parent.index < 0 {
   238  				hash, _ := node.Val.cache()
   239  				state := &nodeIteratorState{
   240  					hash:    common.BytesToHash(hash),
   241  					node:    node.Val,
   242  					parent:  ancestor,
   243  					index:   -1,
   244  					pathlen: len(it.path),
   245  				}
   246  				var path []byte
   247  				if hasTerm(node.Key) {
   248  					path = append(it.path, node.Key[:len(node.Key)-1]...)
   249  				} else {
   250  					path = append(it.path, node.Key...)
   251  				}
   252  				return state, &parent.index, path, nil
   253  			}
   254  		} else if hash, ok := parent.node.(hashNode); ok {
   255  			// Hash node, resolve the hash child from the database
   256  			if parent.index < 0 {
   257  				node, err := it.trie.resolveHash(hash, nil, nil)
   258  				if err != nil {
   259  					return it.stack[len(it.stack)-1], &parent.index, it.path, err
   260  				}
   261  				state := &nodeIteratorState{
   262  					hash:    common.BytesToHash(hash),
   263  					node:    node,
   264  					parent:  ancestor,
   265  					index:   -1,
   266  					pathlen: len(it.path),
   267  				}
   268  				return state, &parent.index, it.path, nil
   269  			}
   270  		}
   271  		// No more child nodes, move back up.
   272  		it.pop()
   273  	}
   274  }
   275  
   276  func (it *nodeIterator) push(state *nodeIteratorState, parentIndex *int, path []byte) {
   277  	it.path = path
   278  	it.stack = append(it.stack, state)
   279  	if parentIndex != nil {
   280  		*parentIndex += 1
   281  	}
   282  }
   283  
   284  func (it *nodeIterator) pop() {
   285  	parent := it.stack[len(it.stack)-1]
   286  	it.path = it.path[:parent.pathlen]
   287  	it.stack = it.stack[:len(it.stack)-1]
   288  }
   289  
   290  func compareNodes(a, b NodeIterator) int {
   291  	cmp := bytes.Compare(a.Path(), b.Path())
   292  	if cmp != 0 {
   293  		return cmp
   294  	}
   295  
   296  	if a.Leaf() && !b.Leaf() {
   297  		return -1
   298  	} else if b.Leaf() && !a.Leaf() {
   299  		return 1
   300  	}
   301  
   302  	cmp = bytes.Compare(a.Hash().Bytes(), b.Hash().Bytes())
   303  	if cmp != 0 {
   304  		return cmp
   305  	}
   306  
   307  	return bytes.Compare(a.LeafBlob(), b.LeafBlob())
   308  }
   309  
   310  type differenceIterator struct {
   311  	a, b  NodeIterator // Nodes returned are those in b - a.
   312  	eof   bool         // Indicates a has run out of elements
   313  	count int          // Number of nodes scanned on either trie
   314  }
   315  
   316  // NewDifferenceIterator constructs a NodeIterator that iterates over elements in b that
   317  // are not in a. Returns the iterator, and a pointer to an integer recording the number
   318  // of nodes seen.
   319  func NewDifferenceIterator(a, b NodeIterator) (NodeIterator, *int) {
   320  	a.Next(true)
   321  	it := &differenceIterator{
   322  		a: a,
   323  		b: b,
   324  	}
   325  	return it, &it.count
   326  }
   327  
   328  func (it *differenceIterator) Hash() common.Hash {
   329  	return it.b.Hash()
   330  }
   331  
   332  func (it *differenceIterator) Parent() common.Hash {
   333  	return it.b.Parent()
   334  }
   335  
   336  func (it *differenceIterator) Leaf() bool {
   337  	return it.b.Leaf()
   338  }
   339  
   340  func (it *differenceIterator) LeafBlob() []byte {
   341  	return it.b.LeafBlob()
   342  }
   343  
   344  func (it *differenceIterator) Path() []byte {
   345  	return it.b.Path()
   346  }
   347  
   348  func (it *differenceIterator) Next(bool) bool {
   349  	// Invariants:
   350  	// - We always advance at least one element in b.
   351  	// - At the start of this function, a's path is lexically greater than b's.
   352  	if !it.b.Next(true) {
   353  		return false
   354  	}
   355  	it.count += 1
   356  
   357  	if it.eof {
   358  		// a has reached eof, so we just return all elements from b
   359  		return true
   360  	}
   361  
   362  	for {
   363  		switch compareNodes(it.a, it.b) {
   364  		case -1:
   365  			// b jumped past a; advance a
   366  			if !it.a.Next(true) {
   367  				it.eof = true
   368  				return true
   369  			}
   370  			it.count += 1
   371  		case 1:
   372  			// b is before a
   373  			return true
   374  		case 0:
   375  			// a and b are identical; skip this whole subtree if the nodes have hashes
   376  			hasHash := it.a.Hash() == common.Hash{}
   377  			if !it.b.Next(hasHash) {
   378  				return false
   379  			}
   380  			it.count += 1
   381  			if !it.a.Next(hasHash) {
   382  				it.eof = true
   383  				return true
   384  			}
   385  			it.count += 1
   386  		}
   387  	}
   388  }
   389  
   390  func (it *differenceIterator) Error() error {
   391  	if err := it.a.Error(); err != nil {
   392  		return err
   393  	}
   394  	return it.b.Error()
   395  }
   396  
   397  type nodeIteratorHeap []NodeIterator
   398  
   399  func (h nodeIteratorHeap) Len() int            { return len(h) }
   400  func (h nodeIteratorHeap) Less(i, j int) bool  { return compareNodes(h[i], h[j]) < 0 }
   401  func (h nodeIteratorHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
   402  func (h *nodeIteratorHeap) Push(x interface{}) { *h = append(*h, x.(NodeIterator)) }
   403  func (h *nodeIteratorHeap) Pop() interface{} {
   404  	n := len(*h)
   405  	x := (*h)[n-1]
   406  	*h = (*h)[0 : n-1]
   407  	return x
   408  }
   409  
   410  type unionIterator struct {
   411  	items *nodeIteratorHeap // Nodes returned are the union of the ones in these iterators
   412  	count int               // Number of nodes scanned across all tries
   413  	err   error             // The error, if one has been encountered
   414  }
   415  
   416  // NewUnionIterator constructs a NodeIterator that iterates over elements in the union
   417  // of the provided NodeIterators. Returns the iterator, and a pointer to an integer
   418  // recording the number of nodes visited.
   419  func NewUnionIterator(iters []NodeIterator) (NodeIterator, *int) {
   420  	h := make(nodeIteratorHeap, len(iters))
   421  	copy(h, iters)
   422  	heap.Init(&h)
   423  
   424  	ui := &unionIterator{
   425  		items: &h,
   426  	}
   427  	return ui, &ui.count
   428  }
   429  
   430  func (it *unionIterator) Hash() common.Hash {
   431  	return (*it.items)[0].Hash()
   432  }
   433  
   434  func (it *unionIterator) Parent() common.Hash {
   435  	return (*it.items)[0].Parent()
   436  }
   437  
   438  func (it *unionIterator) Leaf() bool {
   439  	return (*it.items)[0].Leaf()
   440  }
   441  
   442  func (it *unionIterator) LeafBlob() []byte {
   443  	return (*it.items)[0].LeafBlob()
   444  }
   445  
   446  func (it *unionIterator) Path() []byte {
   447  	return (*it.items)[0].Path()
   448  }
   449  
   450  // Next returns the next node in the union of tries being iterated over.
   451  //
   452  // It does this by maintaining a heap of iterators, sorted by the iteration
   453  // order of their next elements, with one entry for each source trie. Each
   454  // time Next() is called, it takes the least element from the heap to return,
   455  // advancing any other iterators that also point to that same element. These
   456  // iterators are called with descend=false, since we know that any nodes under
   457  // these nodes will also be duplicates, found in the currently selected iterator.
   458  // Whenever an iterator is advanced, it is pushed back into the heap if it still
   459  // has elements remaining.
   460  //
   461  // In the case that descend=false - eg, we're asked to ignore all subnodes of the
   462  // current node - we also advance any iterators in the heap that have the current
   463  // path as a prefix.
   464  func (it *unionIterator) Next(descend bool) bool {
   465  	if len(*it.items) == 0 {
   466  		return false
   467  	}
   468  
   469  	// Get the next key from the union
   470  	least := heap.Pop(it.items).(NodeIterator)
   471  
   472  	// Skip over other nodes as long as they're identical, or, if we're not descending, as
   473  	// long as they have the same prefix as the current node.
   474  	for len(*it.items) > 0 && ((!descend && bytes.HasPrefix((*it.items)[0].Path(), least.Path())) || compareNodes(least, (*it.items)[0]) == 0) {
   475  		skipped := heap.Pop(it.items).(NodeIterator)
   476  		// Skip the whole subtree if the nodes have hashes; otherwise just skip this node
   477  		if skipped.Next(skipped.Hash() == common.Hash{}) {
   478  			it.count += 1
   479  			// If there are more elements, push the iterator back on the heap
   480  			heap.Push(it.items, skipped)
   481  		}
   482  	}
   483  
   484  	if least.Next(descend) {
   485  		it.count += 1
   486  		heap.Push(it.items, least)
   487  	}
   488  
   489  	return len(*it.items) > 0
   490  }
   491  
   492  func (it *unionIterator) Error() error {
   493  	for i := 0; i < len(*it.items); i++ {
   494  		if err := (*it.items)[i].Error(); err != nil {
   495  			return err
   496  		}
   497  	}
   498  	return nil
   499  }