Fix SeekLowerBound where tree contains keys that are prefixes of each other

iradix_test.go

@@ -1533,7 +1533,9 @@ func TestLenTxn(t *testing.T) {
 }
 
 func TestIterateLowerBound(t *testing.T) {
-	fixedLenKeys := []string{
+
+	// these should be defined in order
+	var fixedLenKeys = []string{
 		"00000",
 		"00001",
 		"00004",
@@ -1542,10 +1544,12 @@ func TestIterateLowerBound(t *testing.T) {
 		"20020",
 	}
 
-	mixedLenKeys := []string{
+	// these should be defined in order
+	var mixedLenKeys = []string{
 		"a1",
 		"abc",
 		"barbazboo",
+		"f",
 		"foo",
 		"found",
 		"zap",
@@ -1562,7 +1566,8 @@ func TestIterateLowerBound(t *testing.T) {
 			fixedLenKeys,
 			"00000",
 			fixedLenKeys,
-		}, {
+		},
+		{
 			fixedLenKeys,
 			"00003",
 			[]string{
@@ -1571,72 +1576,86 @@ func TestIterateLowerBound(t *testing.T) {
 				"00020",
 				"20020",
 			},
-		}, {
+		},
+		{
 			fixedLenKeys,
 			"00010",
 			[]string{
 				"00010",
 				"00020",
 				"20020",
 			},
-		}, {
+		},
+		{
 			fixedLenKeys,
 			"20000",
 			[]string{
 				"20020",
 			},
-		}, {
+		},
+		{
 			fixedLenKeys,
 			"20020",
 			[]string{
 				"20020",
 			},
-		}, {
+		},
+		{
 			fixedLenKeys,
 			"20022",
 			[]string{},
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"A", // before all lower case letters
 			mixedLenKeys,
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"a1",
 			mixedLenKeys,
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"b",
 			[]string{
 				"barbazboo",
+				"f",
 				"foo",
 				"found",
 				"zap",
 				"zip",
 			},
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"bar",
 			[]string{
 				"barbazboo",
+				"f",
 				"foo",
 				"found",
 				"zap",
 				"zip",
 			},
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"barbazboo0",
 			[]string{
+				"f",
 				"foo",
 				"found",
 				"zap",
 				"zip",
 			},
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"zippy",
 			[]string{},
-		}, {
+		},
+		{
 			mixedLenKeys,
 			"zi",
 			[]string{
@@ -1665,6 +1684,47 @@ func TestIterateLowerBound(t *testing.T) {
 			"cbacb",
 			[]string{"cbbaa", "cbbab", "cbbbc", "cbcbb", "cbcbc", "cbcca", "ccaaa", "ccabc", "ccaca", "ccacc", "ccbac", "cccaa", "cccac", "cccca"},
 		},
+
+		// Panic case found be TestIterateLowerBoundFuzz.
+		{
+			[]string{"gcgc"},
+			"",
+			[]string{"gcgc"},
+		},
+
+		// We SHOULD support keys that are prefixes of each other despite some
+		// confusion in the original implementation.
+		{
+			[]string{"f", "fo", "foo", "food", "bug"},
+			"foo",
+			[]string{"foo", "food"},
+		},
+
+		// We also support the empty key (which is a prefix of every other key) as a
+		// valid key to insert and search for.
+		{
+			[]string{"f", "fo", "foo", "food", "bug", ""},
+			"foo",
+			[]string{"foo", "food"},
+		},
+		{
+			[]string{"f", "bug", ""},
+			"",
+			[]string{"", "bug", "f"},
+		},
+		{
+			[]string{"f", "bug", "xylophone"},
+			"",
+			[]string{"bug", "f", "xylophone"},
+		},
+
+		// This is a case we realized we were not covering while fixing
+		// SeekReverseLowerBound and could panic before.
+		{
+			[]string{"bar", "foo00", "foo11"},
+			"foo",
+			[]string{"foo00", "foo11"},
+		},
 	}
 
 	for idx, test := range cases {
@@ -1685,9 +1745,7 @@ func TestIterateLowerBound(t *testing.T) {
 			// Get and seek iterator
 			root := r.Root()
 			iter := root.Iterator()
-			if test.search != "" {
-				iter.SeekLowerBound([]byte(test.search))
-			}
+			iter.SeekLowerBound([]byte(test.search))
 
 			// Consume all the keys
 			out := []string{}
@@ -1710,8 +1768,12 @@ type readableString string
 
 func (s readableString) Generate(rand *rand.Rand, size int) reflect.Value {
 	// Pick a random string from a limited alphabet that makes it easy to read the
-	// failure cases. Also never includes a null byte as we don't support that.
-	const letters = "abcdefghijklmnopqrstuvwxyz/-_0123456789"
+	// failure cases.
+	const letters = "abcdefg"
+
+	// Ignore size and make them all shortish to provoke bigger chance of hitting
+	// prefixes and more intersting tree shapes.
+	size = rand.Intn(8)
 
 	b := make([]byte, size)
 	for i := range b {
@@ -1725,17 +1787,16 @@ func TestIterateLowerBoundFuzz(t *testing.T) {
 	set := []string{}
 
 	// This specifies a property where each call adds a new random key to the radix
-	// tree (with a null byte appended since our tree doesn't support one key
-	// being a prefix of another and treats null bytes specially).
+	// tree.
 	//
 	// It also maintains a plain sorted list of the same set of keys and asserts
 	// that iterating from some random key to the end using LowerBound produces
 	// the same list as filtering all sorted keys that are lower.
 
 	radixAddAndScan := func(newKey, searchKey readableString) []string {
-		// Append a null byte
-		key := []byte(newKey + "\x00")
-		r, _, _ = r.Insert(key, nil)
+		r, _, _ = r.Insert([]byte(newKey), nil)
+
+		t.Logf("NewKey: %q, SearchKey: %q", newKey, searchKey)
 
 		// Now iterate the tree from searchKey to the end
 		it := r.Root().Iterator()
@@ -1746,8 +1807,7 @@ func TestIterateLowerBoundFuzz(t *testing.T) {
 			if !ok {
 				break
 			}
-			// Strip the null byte and append to result set
-			result = append(result, string(key[0:len(key)-1]))
+			result = append(result, string(key))
 		}
 		return result
 	}
@@ -1758,14 +1818,19 @@ func TestIterateLowerBoundFuzz(t *testing.T) {
 		sort.Strings(set)
 
 		t.Logf("Current Set: %#v", set)
+		t.Logf("Search Key: %#v %v", searchKey, "" >= string(searchKey))
 
 		result := []string{}
-		var prev string
-		for _, k := range set {
-			if k >= string(searchKey) && k != prev {
+		for i, k := range set {
+			// Check this is not a duplicate of the previous value. Note we don't just
+			// store the last string to compare because empty string is a valid value
+			// in the set and makes comparing on the first iteration awkward.
+			if i > 0 && set[i-1] == k {
+				continue
+			}
+			if k >= string(searchKey) {
 				result = append(result, k)
 			}
-			prev = k
 		}
 		return result
 	}

iter.go

@@ -20,7 +20,7 @@ func (i *Iterator) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) {
 	watch = n.mutateCh
 	search := prefix
 	for {
-		// Check for key exhaution
+		// Check for key exhaustion
 		if len(search) == 0 {
 			i.node = n
 			return
@@ -60,10 +60,13 @@ func (i *Iterator) recurseMin(n *Node) *Node {
 	if n.leaf != nil {
 		return n
 	}
-	if len(n.edges) > 0 {
+	nEdges := len(n.edges)
+	if nEdges > 1 {
 		// Add all the other edges to the stack (the min node will be added as
 		// we recurse)
 		i.stack = append(i.stack, n.edges[1:])
+	}
+	if nEdges > 0 {
 		return i.recurseMin(n.edges[0].node)
 	}
 	// Shouldn't be possible
@@ -77,16 +80,32 @@ func (i *Iterator) recurseMin(n *Node) *Node {
 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.
+	// leaf with the lower bound. Note that the iterator will still recurse into
+	// children that we don't traverse on the way to the reverse lower bound as it
+	// walks the stack.
 	i.stack = []edges{}
+	// i.node starts off in the common case as pointing to the root node of the
+	// tree. By the time we return we have either found a lower bound and setup
+	// the stack to traverse all larger keys, or we have not and the stack and
+	// node should both be nil to prevent the iterator from assuming it is just
+	// iterating the whole tree from the root node. Either way this needs to end
+	// up as nil so just set it here.
 	n := i.node
+	i.node = nil
 	search := key
 
 	found := func(n *Node) {
-		i.node = n
 		i.stack = append(i.stack, edges{edge{node: n}})
 	}
 
+	findMin := func(n *Node) {
+		n = i.recurseMin(n)
+		if n != nil {
+			found(n)
+			return
+		}
+	}
+
 	for {
 		// Compare current prefix with the search key's same-length prefix.
 		var prefixCmp int
@@ -100,10 +119,7 @@ func (i *Iterator) SeekLowerBound(key []byte) {
 			// Prefix is larger, that means the lower bound is greater than the search
 			// and from now on we need to follow the minimum path to the smallest
 			// leaf under this subtree.
-			n = i.recurseMin(n)
-			if n != nil {
-				found(n)
-			}
+			findMin(n)
 			return
 		}
 
@@ -115,27 +131,29 @@ func (i *Iterator) SeekLowerBound(key []byte) {
 		}
 
 		// Prefix is equal, we are still heading for an exact match. If this is a
-		// leaf we're done.
-		if n.leaf != nil {
-			if bytes.Compare(n.leaf.key, key) < 0 {
-				i.node = nil
-				return
-			}
+		// leaf and an exact match we're done.
+		if n.leaf != nil && bytes.Equal(n.leaf.key, key) {
 			found(n)
 			return
 		}
 
-		// Consume the search prefix
-		if len(n.prefix) > len(search) {
-			search = []byte{}
-		} else {
-			search = search[len(n.prefix):]
+		// Consume the search prefix if the current node has one. Note that this is
+		// safe because if n.prefix is longer than the search slice prefixCmp would
+		// have been > 0 above and the method would have already returned.
+		search = search[len(n.prefix):]
+
+		if len(search) == 0 {
+			// We've exhausted the search key, but the current node is not an exact
+			// match or not a leaf. That means that the leaf value if it exists, and
+			// all child nodes must be strictly greater, the smallest key in this
+			// subtree must be the lower bound.
+			findMin(n)
+			return
 		}
 
 		// Otherwise, take the lower bound next edge.
 		idx, lbNode := n.getLowerBoundEdge(search[0])
 		if lbNode == nil {
-			i.node = nil
 			return
 		}
 
@@ -144,7 +162,6 @@ func (i *Iterator) SeekLowerBound(key []byte) {
 			i.stack = append(i.stack, n.edges[idx+1:])
 		}
 
-		i.node = lbNode
 		// Recurse
 		n = lbNode
 	}