Basic working SeekLowerBound to allow range scans

iradix_test.go

@@ -1529,3 +1529,184 @@ func TestLenTxn(t *testing.T) {
 		t.Fatalf("tree len should be zero, got %d", r.Len())
 	}
 }
+
+func TestIterateLowerBound(t *testing.T) {
+	r := New()
+
+	keys := []string{
+		"00000",
+		"00001",
+		"00004",
+		"00010",
+		"00020",
+		"20020",
+	}
+	for _, k := range keys {
+		r, _, _ = r.Insert([]byte(k), nil)
+	}
+	if r.Len() != len(keys) {
+		t.Fatalf("bad len: %v %v", r.Len(), len(keys))
+	}
+
+	type exp struct {
+		inp string
+		out []string
+	}
+	cases := []exp{
+		exp{
+			"00000",
+			keys,
+		},
+		exp{
+			"00003",
+			[]string{
+				"00004",
+				"00010",
+				"00020",
+				"20020",
+			},
+		},
+		exp{
+			"00010",
+			[]string{
+				"00010",
+				"00020",
+				"20020",
+			},
+		},
+		exp{
+			"20000",
+			[]string{
+				"20020",
+			},
+		},
+		exp{
+			"20020",
+			[]string{
+				"20020",
+			},
+		},
+		exp{
+			"20022",
+			[]string{},
+		},
+	}
+
+	root := r.Root()
+	for idx, test := range cases {
+		iter := root.Iterator()
+		if test.inp != "" {
+			iter.SeekLowerBound([]byte(test.inp))
+		}
+
+		// Consume all the keys
+		out := []string{}
+		for {
+			key, _, ok := iter.Next()
+			if !ok {
+				break
+			}
+			out = append(out, string(key))
+		}
+		if !reflect.DeepEqual(out, test.out) {
+			t.Fatalf("mis-match[%d]: key=%s\n  got=%v\n  want=%v", idx, test.inp,
+				out, test.out)
+		}
+	}
+}
+
+func TestIterateLowerBoundMixedLenKeys(t *testing.T) {
+	r := New()
+
+	keys := []string{
+		"a1",
+		"abc",
+		"barbazboo",
+		"foo",
+		"found",
+		"zap",
+		"zip",
+	}
+	for _, k := range keys {
+		r, _, _ = r.Insert([]byte(k), nil)
+	}
+	if r.Len() != len(keys) {
+		t.Fatalf("bad len: %v %v", r.Len(), len(keys))
+	}
+
+	type exp struct {
+		inp string
+		out []string
+	}
+	cases := []exp{
+		exp{
+			"A", // before all lower case letters
+			keys,
+		},
+		exp{
+			"a1",
+			keys,
+		},
+		exp{
+			"b",
+			[]string{
+				"barbazboo",
+				"foo",
+				"found",
+				"zap",
+				"zip",
+			},
+		},
+		exp{
+			"bar",
+			[]string{
+				"barbazboo",
+				"foo",
+				"found",
+				"zap",
+				"zip",
+			},
+		},
+		exp{
+			"barbazboo0",
+			[]string{
+				"foo",
+				"found",
+				"zap",
+				"zip",
+			},
+		},
+		exp{
+			"zippy",
+			[]string{},
+		},
+		exp{
+			"zi",
+			[]string{
+				"zip",
+			},
+		},
+	}
+
+	root := r.Root()
+	for idx, test := range cases {
+		iter := root.Iterator()
+		if test.inp != "" {
+			iter.SeekLowerBound([]byte(test.inp))
+		}
+
+		// Consume all the keys
+		out := []string{}
+		for {
+			key, _, ok := iter.Next()
+			if !ok {
+				break
+			}
+			out = append(out, string(key))
+		}
+		if !reflect.DeepEqual(out, test.out) {
+			t.Fatalf("mis-match[%d]: key=%s\n  got=%v\n  want=%v", idx, test.inp,
+				out, test.out)
+		}
+	}
+}

iter.go

@@ -1,6 +1,8 @@
 package iradix
 
-import "bytes"
+import (
+	"bytes"
+)
 
 // Iterator is used to iterate over a set of nodes
 // in pre-order
@@ -53,6 +55,69 @@ func (i *Iterator) SeekPrefix(prefix []byte) {
 	i.SeekPrefixWatch(prefix)
 }
 
+// SeekLowerBound is used to seek the iterator to the smallest key that is
+// greater or equal to the given key. There is no watch variant as it's hard to
+// predict based on the radix structure which node(s) changes might affect the
+// result.
+func (i *Iterator) SeekLowerBound(key []byte) {
+	// Wipe the stack. Unlike Prefix iteration, we need to build the stack as we
+	// go because we need only a subset of edges of many nodes in the path to the
+	// leaf with the lower bound.
+	i.stack = []edges{}
+	n := i.node
+	search := key
+
+	found := func(n *Node) {
+		i.node = n
+		i.stack = append(i.stack, edges{edge{node: n}})
+	}
+
+	for {
+		// Consume the search prefix
+		if len(n.prefix) > len(search) {
+			search = []byte{}
+		} else {
+			search = search[len(n.prefix):]
+		}
+
+		// Is this a leaf? If so check if it's lower than the key (no lower bound
+		// exists).
+		if n.leaf != nil {
+			if bytes.Compare(n.leaf.key, key) < 0 {
+				i.node = nil
+				return
+			}
+
+			// We are a leaf that is greater or equal, that means we are the lower
+			// bound.
+			found(n)
+			return
+		}
+
+		// Check for key exhaustion
+		if len(search) == 0 {
+			found(n)
+			return
+		}
+
+		// Not a leaf, look for the lower bound edge
+		idx, lbNode := n.getLowerBoundEdge(search[0])
+		if lbNode == nil {
+			i.node = nil
+			return
+		}
+
+		// Create stack edges for the all strictly higher edges in this node.
+		if idx+1 < len(n.edges) {
+			i.stack = append(i.stack, n.edges[idx+1:])
+		}
+
+		i.node = lbNode
+		// Recurse
+		n = lbNode
+	}
+}
+
 // Next returns the next node in order
 func (i *Iterator) Next() ([]byte, interface{}, bool) {
 	// Initialize our stack if needed

node.go

@@ -79,6 +79,18 @@ func (n *Node) getEdge(label byte) (int, *Node) {
 	return -1, nil
 }
 
+func (n *Node) getLowerBoundEdge(label byte) (int, *Node) {
+	num := len(n.edges)
+	idx := sort.Search(num, func(i int) bool {
+		return n.edges[i].label >= label
+	})
+	// we want lower bound behavior so return even if it's not an exact match
+	if idx < num {
+		return idx, n.edges[idx].node
+	}
+	return -1, nil
+}
+
 func (n *Node) delEdge(label byte) {
 	num := len(n.edges)
 	idx := sort.Search(num, func(i int) bool {