add regular seek

tonistiigi · tonistiigi · commit 826af9ccf0fe · 2017-08-03T11:56:27.000-07:00
Signed-off-by: Tonis Tiigi &lt;tonistiigi@gmail.com&gt;
diff --git a/node.go b/node.go
@@ -274,6 +274,66 @@ func (n *Node) WalkPath(path []byte, fn WalkFn) {
 	}
 }
 
+func (n *Node) Seek(prefix []byte) *Seeker {
+	search := prefix
+	p := &pos{n: n}
+	for {
+		// Check for key exhaution
+		if len(search) == 0 {
+			return &Seeker{p}
+		}
+
+		num := len(n.edges)
+		idx := sort.Search(num, func(i int) bool {
+			return n.edges[i].label >= search[0]
+		})
+		p.current = idx
+		if idx < len(n.edges) {
+			n = n.edges[idx].node
+			if bytes.HasPrefix(search, n.prefix) && len(n.edges) > 0 {
+				search = search[len(n.prefix):]
+				p.current++
+				p = &pos{n: n, prev: p}
+				continue
+			}
+		}
+		p.current++
+		return &Seeker{p}
+	}
+}
+
+type Seeker struct {
+	*pos
+}
+
+type pos struct {
+	n       *Node
+	current int
+	prev    *pos
+	isLeaf  bool
+}
+
+func (s *Seeker) Next() (k []byte, v interface{}, ok bool) {
+	if s.current >= len(s.n.edges) {
+		if s.prev == nil {
+			return nil, nil, false
+		}
+		s.pos = s.prev
+		return s.Next()
+	}
+
+	edge := s.n.edges[s.current]
+	s.current++
+	if edge.node.leaf != nil && !s.isLeaf {
+		s.isLeaf = true
+		s.current--
+		return edge.node.leaf.key, edge.node.leaf.val, true
+	}
+	s.isLeaf = false
+	s.pos = &pos{n: edge.node, prev: s.pos}
+	return s.Next()
+}
+
 // recursiveWalk is used to do a pre-order walk of a node
 // recursively. Returns true if the walk should be aborted
 func recursiveWalk(n *Node, fn WalkFn) bool {
