github.com/alexandrestein/gods@v1.0.1/trees/avltree/iterator.go (about)

     1  // Copyright (c) 2017, Benjamin Scher Purcell. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package avltree
     6  
     7  import "github.com/alexandrestein/gods/containers"
     8  
     9  func assertIteratorImplementation() {
    10  	var _ containers.ReverseIteratorWithKey = (*Iterator)(nil)
    11  }
    12  
    13  // Iterator holding the iterator's state
    14  type Iterator struct {
    15  	tree     *Tree
    16  	node     *Node
    17  	position position
    18  }
    19  
    20  type position byte
    21  
    22  const (
    23  	begin, between, end position = 0, 1, 2
    24  )
    25  
    26  // Iterator returns a stateful iterator whose elements are key/value pairs.
    27  func (tree *Tree) Iterator() containers.ReverseIteratorWithKey {
    28  	return &Iterator{tree: tree, node: nil, position: begin}
    29  }
    30  
    31  // Next moves the iterator to the next element and returns true if there was a next element in the container.
    32  // If Next() returns true, then next element's key and value can be retrieved by Key() and Value().
    33  // If Next() was called for the first time, then it will point the iterator to the first element if it exists.
    34  // Modifies the state of the iterator.
    35  func (iterator *Iterator) Next() bool {
    36  	switch iterator.position {
    37  	case begin:
    38  		iterator.position = between
    39  		iterator.node = iterator.tree.Left()
    40  	case between:
    41  		iterator.node = iterator.node.Next()
    42  	}
    43  
    44  	if iterator.node == nil {
    45  		iterator.position = end
    46  		return false
    47  	}
    48  	return true
    49  }
    50  
    51  // Prev moves the iterator to the next element and returns true if there was a previous element in the container.
    52  // If Prev() returns true, then next element's key and value can be retrieved by Key() and Value().
    53  // If Prev() was called for the first time, then it will point the iterator to the first element if it exists.
    54  // Modifies the state of the iterator.
    55  func (iterator *Iterator) Prev() bool {
    56  	switch iterator.position {
    57  	case end:
    58  		iterator.position = between
    59  		iterator.node = iterator.tree.Right()
    60  	case between:
    61  		iterator.node = iterator.node.Prev()
    62  	}
    63  
    64  	if iterator.node == nil {
    65  		iterator.position = begin
    66  		return false
    67  	}
    68  	return true
    69  }
    70  
    71  // Value returns the current element's value.
    72  // Does not modify the state of the iterator.
    73  func (iterator *Iterator) Value() interface{} {
    74  	if iterator.node == nil {
    75  		return nil
    76  	}
    77  	return iterator.node.Value
    78  }
    79  
    80  // Key returns the current element's key.
    81  // Does not modify the state of the iterator.
    82  func (iterator *Iterator) Key() interface{} {
    83  	if iterator.node == nil {
    84  		return nil
    85  	}
    86  	return iterator.node.Key
    87  }
    88  
    89  // Begin resets the iterator to its initial state (one-before-first)
    90  // Call Next() to fetch the first element if any.
    91  func (iterator *Iterator) Begin() {
    92  	iterator.node = nil
    93  	iterator.position = begin
    94  }
    95  
    96  // End moves the iterator past the last element (one-past-the-end).
    97  // Call Prev() to fetch the last element if any.
    98  func (iterator *Iterator) End() {
    99  	iterator.node = nil
   100  	iterator.position = end
   101  }
   102  
   103  // First moves the iterator to the first element and returns true if there was a first element in the container.
   104  // If First() returns true, then first element's key and value can be retrieved by Key() and Value().
   105  // Modifies the state of the iterator
   106  func (iterator *Iterator) First() bool {
   107  	iterator.Begin()
   108  	return iterator.Next()
   109  }
   110  
   111  // Last moves the iterator to the last element and returns true if there was a last element in the container.
   112  // If Last() returns true, then last element's key and value can be retrieved by Key() and Value().
   113  // Modifies the state of the iterator.
   114  func (iterator *Iterator) Last() bool {
   115  	iterator.End()
   116  	return iterator.Prev()
   117  }