Have ability to store depth on tree traversal

zmanian · Feb 1, 2018 · 86d53c9 · 86d53c9
1 parent 6426fca
commit 86d53c9
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 13 deletions.
diff --git a/node.go b/node.go
@@ -505,10 +505,16 @@ func (node *Node) balance(t *Tree) (newSelf *Node, orphaned []*Node) {
 
 // traverse is a wrapper over traverseInRange when we want the whole tree
 func (node *Node) traverse(t *Tree, ascending bool, cb func(*Node) bool) bool {
-	return node.traverseInRange(t, nil, nil, ascending, false, cb)
+	return node.traverseInRange(t, nil, nil, ascending, false, 0, func(node *Node, depth uint8) bool {
+		return cb(node)
+	})
 }
 
-func (node *Node) traverseInRange(t *Tree, start, end []byte, ascending bool, inclusive bool, cb func(*Node) bool) bool {
+func (node *Node) traverseWithDepth(t *Tree, ascending bool, cb func(*Node, uint8) bool) bool {
+	return node.traverseInRange(t, nil, nil, ascending, false, 0, cb)
+}
+
+func (node *Node) traverseInRange(t *Tree, start, end []byte, ascending bool, inclusive bool, depth uint8, cb func(*Node, uint8) bool) bool {
 	afterStart := start == nil || bytes.Compare(start, node.key) <= 0
 	beforeEnd := end == nil || bytes.Compare(node.key, end) < 0
 	if inclusive {
@@ -518,7 +524,7 @@ func (node *Node) traverseInRange(t *Tree, start, end []byte, ascending bool, in
 	stop := false
 	if afterStart && beforeEnd {
 		// IterateRange ignores this if not leaf
-		stop = cb(node)
+		stop = cb(node, depth)
 	}
 	if stop {
 		return stop
@@ -530,24 +536,24 @@ func (node *Node) traverseInRange(t *Tree, start, end []byte, ascending bool, in
 	if ascending {
 		// check lower nodes, then higher
 		if afterStart {
-			stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, cb)
+			stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, cb)
 		}
 		if stop {
 			return stop
 		}
 		if beforeEnd {
-			stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, cb)
+			stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, cb)
 		}
 	} else {
 		// check the higher nodes first
 		if beforeEnd {
-			stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, cb)
+			stop = node.getRightNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, cb)
 		}
 		if stop {
 			return stop
 		}
 		if afterStart {
-			stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, cb)
+			stop = node.getLeftNode(t).traverseInRange(t, start, end, ascending, inclusive, depth+1, cb)
 		}
 	}
 

diff --git a/serialize.go b/serialize.go
@@ -1,9 +1,15 @@
 package iavl
 
-// NodeData groups together a key and a value for return codes.
+import (
+	"fmt"
+	"math"
+)
+
+// NodeData groups together a key, value and depth.
 type NodeData struct {
 	Key   []byte
 	Value []byte
+	Depth uint8
 }
 
 // SerializeFunc is any implementation that can serialize
@@ -24,9 +30,9 @@ func Restore(empty *Tree, kvs []NodeData) {
 // when recovering, it will create a different.
 func InOrderSerialize(t *Tree, root *Node) []NodeData {
 	res := make([]NodeData, 0, root.size)
-	root.traverse(t, true, func(node *Node) bool {
+	root.traverseWithDepth(t, true, func(node *Node, depth uint8) bool {
 		if node.height == 0 {
-			kv := NodeData{Key: node.key, Value: node.value}
+			kv := NodeData{Key: node.key, Value: node.value, Depth: depth}
 			res = append(res, kv)
 		}
 		return false
@@ -67,7 +73,7 @@ func StableSerializeBFS(t *Tree, root *Node) []NodeData {
 
 	nds := make([]NodeData, size)
 	for i, k := range keys {
-		nds[i] = NodeData{k, visited[string(k)]}
+		nds[i] = NodeData{k, visited[string(k)], 0}
 	}
 	return nds
 }

diff --git a/tree.go b/tree.go
@@ -222,7 +222,7 @@ func (t *Tree) IterateRange(start, end []byte, ascending bool, fn func(key []byt
 	if t.root == nil {
 		return false
 	}
-	return t.root.traverseInRange(t, start, end, ascending, false, func(node *Node) bool {
+	return t.root.traverseInRange(t, start, end, ascending, false, 0, func(node *Node, _ uint8) bool {
 		if node.height == 0 {
 			return fn(node.key, node.value)
 		} else {
@@ -237,7 +237,7 @@ func (t *Tree) IterateRangeInclusive(start, end []byte, ascending bool, fn func(
 	if t.root == nil {
 		return false
 	}
-	return t.root.traverseInRange(t, start, end, ascending, true, func(node *Node) bool {
+	return t.root.traverseInRange(t, start, end, ascending, true, 0, func(node *Node, _ uint8) bool {
 		if node.height == 0 {
 			return fn(node.key, node.value)
 		} else {