IAVL Iterator

immutable_tree.go

@@ -223,3 +223,71 @@ func (t *ImmutableTree) nodeSize() int {
 	})
 	return size
 }
+
+type Iterator struct {
+	start, end []byte
+
+	key, value []byte
+
+	valid bool
+
+	// next is the argument for retrieving the sequence starting from the next element.
+	// By calling next manually, Iterator can control the execution flow.
+	next traversal
+}
+
+func (t *ImmutableTree) Iterator(start, end []byte, ascending bool) *Iterator {
+	iter := &Iterator{
+		start: start,
+		end:   end,
+		valid: true,
+		next:  t.root.traversal(t, start, end, ascending, false, 0, false, true, nil),
+	}
+
+	iter.Next()
+	return iter
+}
+
+var _ dbm.Iterator = &Iterator{}
+
+func (iter *Iterator) Domain() ([]byte, []byte) {
+	return iter.start, iter.end
+}
+
+func (iter *Iterator) Valid() bool {
+	return iter.valid
+}
+
+func (iter *Iterator) Key() []byte {
+	return iter.key
+}
+
+func (iter *Iterator) Value() []byte {
+	return iter.value
+}
+
+func (iter *Iterator) Next() {
+	if iter.next == nil {
+		iter.valid = false
+		return
+	}
+	node, _, next := iter.next()
+	iter.next = next
+
+	if node.height == 0 {
+		iter.key, iter.value = node.key, node.value
+		return
+	}
+
+	iter.Next()
+}
+
+func (iter *Iterator) Close() error {
+	iter.next = nil
+	iter.valid = false
+	return nil
+}
+
+func (iter *Iterator) Error() error {
+	return nil
+}

node.go

@@ -452,63 +452,106 @@ func (node *Node) traversePost(t *ImmutableTree, ascending bool, cb func(*Node)
 	})
 }
 
+type traversal func() (*Node, uint8, traversal)
+
 func (node *Node) traverseInRange(t *ImmutableTree, start, end []byte, ascending bool, inclusive bool, depth uint8, post bool, cb func(*Node, uint8) bool) bool {
+	stop := false
+	next := node.traversal(t, start, end, ascending, inclusive, depth, post, true, nil)
+	var node2 *Node
+	var depth2 uint8
+	for next != nil {
+		node2, depth2, next = next()
+		stop = cb(node2, depth2)
+		if stop {
+			return stop
+		}
+	}
+	return stop
+}
+
+// `traversal` returns the delayed execution of recursive traveral in order to give the caller the control of the execution flow.
+// The caller can either call the next traversal to proceed, or simply discard the function variable to stop iteration.
+// The condition is to determine whether this traversal is required or not, depending on the caller is included in the search
+// range or not.
+func (node *Node) traversal(
+	t *ImmutableTree,
+	start, end []byte, ascending bool, inclusive bool,
+	depth uint8,
+	post bool,
+	condition bool,
+	continuation traversal,
+) traversal {
 	if node == nil {
-		return false
+		return continuation
+	}
+
+	if !condition {
+		return continuation
 	}
+
 	afterStart := start == nil || bytes.Compare(start, node.key) < 0
 	startOrAfter := start == nil || bytes.Compare(start, node.key) <= 0
 	beforeEnd := end == nil || bytes.Compare(node.key, end) < 0
 	if inclusive {
 		beforeEnd = end == nil || bytes.Compare(node.key, end) <= 0
 	}
 
-	// Run callback per inner/leaf node.
-	stop := false
-	if !post && (!node.isLeaf() || (startOrAfter && beforeEnd)) {
-		stop = cb(node, depth)
-		if stop {
-			return stop
+	inner := func(continuation traversal) traversal {
+		if !node.isLeaf() {
+			if ascending {
+				// if node is a branch node and the order is ascending, then;
+				// return the traversal for the left nodes, which will then proceed on the continuation of;
+				return node.getLeftNode(t).traversal(t, start, end, ascending, inclusive, depth+1, post, afterStart,
+					// the traversal for the right nodes, which will then proceed on the continuation of;
+					node.getRightNode(t).traversal(t, start, end, ascending, inclusive, depth+1, post, beforeEnd,
+						// the delayed traversal for the parent nodes and their children(which is passed as an argument)
+						continuation))
+			} else {
+				// if node is a branch node and the order is not ascending, then;
+				// return the traversal for the right nodes, which will then proceed on the continuation of;
+				return node.getRightNode(t).traversal(t, start, end, ascending, inclusive, depth+1, post, beforeEnd,
+					// the traversal for the left nodes, which will then proceed on the continuation of;
+					node.getLeftNode(t).traversal(t, start, end, ascending, inclusive, depth+1, post, afterStart,
+						// the delayed traversal for the parent nodes and their children(which is passed as an argument)
+						continuation))
+			}
 		}
+		// if node is not a branch node, we don't care about it in this closure.
+		return continuation
 	}
 
-	if !node.isLeaf() {
-		if ascending {
-			// check lower nodes, then higher
-			if afterStart {
-				stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, post, cb)
-			}
-			if stop {
-				return stop
-			}
-			if beforeEnd {
-				stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, post, cb)
-			}
-		} else {
-			// check the higher nodes first
-			if beforeEnd {
-				stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, post, cb)
-			}
-			if stop {
-				return stop
-			}
-			if afterStart {
-				stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, post, cb)
-			}
+	// Run callback per inner/leaf node.
+
+	// case of preorder traversal.
+	if !post && (!node.isLeaf() || (startOrAfter && beforeEnd)) {
+		return func() (*Node, uint8, traversal) {
+			// return the traversal for this node, which will then proceed on the continuation of;
+			return node, depth,
+				// the traversal for the inner nodes, which will then proceed on the continuation of;
+				inner(
+					// the delayed traversal for the parent nodes and their children(which is passed as an argument)
+					continuation,
+				)
 		}
 	}
-	if stop {
-		return stop
-	}
 
+	// case of postorder traversal.
 	if post && (!node.isLeaf() || (startOrAfter && beforeEnd)) {
-		stop = cb(node, depth)
-		if stop {
-			return stop
+		return func() (*Node, uint8, traversal) {
+			// return the traversal for the inner nodes, which will then proceed on the continuation of;
+			return inner(func() (*Node, uint8, traversal) {
+				// the traversal for this node, which will then proceed on the continuation of;
+				return node, depth,
+					// the delayed traversal for the parent nodes and their children(which is passed as an argument)
+					continuation
+			})()
 		}
 	}
 
-	return stop
+	// ignore this node, keep traverse on children and then proceed to delayed continuation.
+	return inner(
+		continuation,
+	)
 }
 
 // Only used in testing...