github.com/iotexproject/iotex-core@v1.14.1-rc1/db/trie/mptrie/leafiterator.go (about)

     1  // Copyright (c) 2018 IoTeX
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package mptrie
     7  
     8  import (
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/iotexproject/iotex-core/db/trie"
    12  )
    13  
    14  // LeafIterator defines an iterator to go through all the leaves under given node
    15  type LeafIterator struct {
    16  	cli   client
    17  	stack []node
    18  }
    19  
    20  // NewLeafIterator returns a new leaf iterator
    21  func NewLeafIterator(tr trie.Trie) (trie.Iterator, error) {
    22  	mpt, ok := tr.(*merklePatriciaTrie)
    23  	if !ok {
    24  		return nil, errors.New("trie is not supported type")
    25  	}
    26  	stack := []node{mpt.root}
    27  
    28  	return &LeafIterator{cli: mpt, stack: stack}, nil
    29  }
    30  
    31  // Next moves iterator to next node
    32  func (li *LeafIterator) Next() ([]byte, []byte, error) {
    33  	for len(li.stack) > 0 {
    34  		size := len(li.stack)
    35  		node := li.stack[size-1]
    36  		li.stack = li.stack[:size-1]
    37  		if hn, ok := node.(*hashNode); ok {
    38  			node, err := hn.LoadNode(li.cli)
    39  			if err != nil {
    40  				return nil, nil, err
    41  			}
    42  			li.stack = append(li.stack, node)
    43  			continue
    44  		}
    45  		if ln, ok := node.(leaf); ok {
    46  			key := ln.Key()
    47  			value := ln.Value()
    48  
    49  			return append(key[:0:0], key...), append(value[:0:0], value...), nil
    50  		}
    51  		if bn, ok := node.(branch); ok {
    52  			children := bn.Children()
    53  			li.stack = append(li.stack, children...)
    54  			continue
    55  		}
    56  		if en, ok := node.(extension); ok {
    57  			li.stack = append(li.stack, en.Child())
    58  			continue
    59  		}
    60  		return nil, nil, errors.New("unexpected node type")
    61  	}
    62  
    63  	return nil, nil, trie.ErrEndOfIterator
    64  }