github.com/jeffallen/go-ethereum@v1.1.4-0.20150910155051-571d3236c49c/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  )
    22  
    23  type Iterator struct {
    24  	trie *Trie
    25  
    26  	Key   []byte
    27  	Value []byte
    28  }
    29  
    30  func NewIterator(trie *Trie) *Iterator {
    31  	return &Iterator{trie: trie, Key: nil}
    32  }
    33  
    34  func (self *Iterator) Next() bool {
    35  	self.trie.mu.Lock()
    36  	defer self.trie.mu.Unlock()
    37  
    38  	isIterStart := false
    39  	if self.Key == nil {
    40  		isIterStart = true
    41  		self.Key = make([]byte, 32)
    42  	}
    43  
    44  	key := RemTerm(CompactHexDecode(self.Key))
    45  	k := self.next(self.trie.root, key, isIterStart)
    46  
    47  	self.Key = []byte(DecodeCompact(k))
    48  
    49  	return len(k) > 0
    50  }
    51  
    52  func (self *Iterator) next(node Node, key []byte, isIterStart bool) []byte {
    53  	if node == nil {
    54  		return nil
    55  	}
    56  
    57  	switch node := node.(type) {
    58  	case *FullNode:
    59  		if len(key) > 0 {
    60  			k := self.next(node.branch(key[0]), key[1:], isIterStart)
    61  			if k != nil {
    62  				return append([]byte{key[0]}, k...)
    63  			}
    64  		}
    65  
    66  		var r byte
    67  		if len(key) > 0 {
    68  			r = key[0] + 1
    69  		}
    70  
    71  		for i := r; i < 16; i++ {
    72  			k := self.key(node.branch(byte(i)))
    73  			if k != nil {
    74  				return append([]byte{i}, k...)
    75  			}
    76  		}
    77  
    78  	case *ShortNode:
    79  		k := RemTerm(node.Key())
    80  		if vnode, ok := node.Value().(*ValueNode); ok {
    81  			switch bytes.Compare([]byte(k), key) {
    82  			case 0:
    83  				if isIterStart {
    84  					self.Value = vnode.Val()
    85  					return k
    86  				}
    87  			case 1:
    88  				self.Value = vnode.Val()
    89  				return k
    90  			}
    91  		} else {
    92  			cnode := node.Value()
    93  
    94  			var ret []byte
    95  			skey := key[len(k):]
    96  			if BeginsWith(key, k) {
    97  				ret = self.next(cnode, skey, isIterStart)
    98  			} else if bytes.Compare(k, key[:len(k)]) > 0 {
    99  				return self.key(node)
   100  			}
   101  
   102  			if ret != nil {
   103  				return append(k, ret...)
   104  			}
   105  		}
   106  	}
   107  
   108  	return nil
   109  }
   110  
   111  func (self *Iterator) key(node Node) []byte {
   112  	switch node := node.(type) {
   113  	case *ShortNode:
   114  		// Leaf node
   115  		if vnode, ok := node.Value().(*ValueNode); ok {
   116  			k := RemTerm(node.Key())
   117  			self.Value = vnode.Val()
   118  
   119  			return k
   120  		} else {
   121  			k := RemTerm(node.Key())
   122  			return append(k, self.key(node.Value())...)
   123  		}
   124  	case *FullNode:
   125  		if node.Value() != nil {
   126  			self.Value = node.Value().(*ValueNode).Val()
   127  
   128  			return []byte{16}
   129  		}
   130  
   131  		for i := 0; i < 16; i++ {
   132  			k := self.key(node.branch(byte(i)))
   133  			if k != nil {
   134  				return append([]byte{byte(i)}, k...)
   135  			}
   136  		}
   137  	}
   138  
   139  	return nil
   140  }