util/interval: implement Iterator method for btree

1. Implement Iterator method for btree and add tests.
2. Restore Range.Equal method and fix the wrong logic of Range.Equal.
   The existing Equal method considers Range A and Range B to be equal
   if A's Start is equal to B's End and A's End is equal to B's Start.
   The correct logic is to consider Range A and Range B to be equal if
   A's Start is equal to B's Start and A's End is equal to B's End. See
   #16079.
3. Add complile time checks to ensure that btree and llrbTree implement
   the Tree interface.
4. s/llbrNode/llrbNode/g and s/llbrTree/llrbTree/g.
5. Move comment from llrbTree to Tree.

pkg/util/interval/btree_based_interval.go

@@ -16,10 +16,6 @@
 //
 // Author: Jingguo Yao (yaojingguo@gmail.com)
 
-// Package interval implements an interval tree based on an augmented B-tree.
-// For all the functions which have a fast argment, fast being true means a fast
-// operation which does not adjust node ranges. If fast is false, node ranges
-// are adjusted.
 package interval
 
 import (
@@ -33,9 +29,11 @@ const (
 	DefaultBTreeMinimumDegree = 32
 )
 
+var _ = newBTree
+
 // newBTree creates a new interval tree with the given overlapper function and
 // the default B-tree minimum degree.
-func newBTree(overlapper Overlapper) btree {
+func newBTree(overlapper Overlapper) *btree {
 	return newBTreeWithDegree(overlapper, DefaultBTreeMinimumDegree)
 }
 
@@ -45,11 +43,11 @@ func newBTree(overlapper Overlapper) btree {
 //
 // newBTreeWithDegree(overlapper, 2), for example, will create a 2-3-4 tree (each
 // node contains 1-3 Interfaces and 2-4 children).
-func newBTreeWithDegree(overlapper Overlapper, minimumDegree int) btree {
+func newBTreeWithDegree(overlapper Overlapper, minimumDegree int) *btree {
 	if minimumDegree < 2 {
 		panic("bad minimum degree")
 	}
-	return btree{
+	return &btree{
 		MinimumDegree: minimumDegree,
 		Overlapper:    overlapper,
 	}
@@ -276,15 +274,10 @@ func (t *btree) isEmpty() bool {
 	return t.root == nil || len(t.root.items) == 0
 }
 
-// Get returns a slice of Interfaces that overlap r in the tree. The slice is
-// sorted nondecreasingly by interval start.
 func (t *btree) Get(r Range) (o []Interface) {
 	return t.GetWithOverlapper(r, t.Overlapper)
 }
 
-// GetWithOverlapper returns a slice of Interfaces that overlap r in the tree
-// using the provided overlapper function. The slice is sorted nondecreasingly
-// by interval start.
 func (t *btree) GetWithOverlapper(r Range, overlapper Overlapper) (o []Interface) {
 	if err := rangeError(r); err != nil {
 		return
@@ -296,11 +289,6 @@ func (t *btree) GetWithOverlapper(r Range, overlapper Overlapper) (o []Interface
 	return
 }
 
-// DoMatching performs fn on all intervals stored in the tree that overlaps r.
-// The traversal is done in the nondecreasing order of interval start. A boolean
-// is returned indicating whether the traversal was interrupted by an Operation
-// returning true. If fn alters stored intervals' sort relationships, future
-// tree operation behaviors are undefined.
 func (t *btree) DoMatching(fn Operation, r Range) bool {
 	if err := rangeError(r); err != nil {
 		return false
@@ -406,11 +394,6 @@ func (n *node) exclusiveDoMatch(fn Operation, r Range, overlapper Overlapper) (d
 	return
 }
 
-// Do performs fn on all intervals stored in the tree. The traversal is done in
-// the nondecreasing order of interval start. A boolean is returned indicating
-// whether the traversal was interrupted by an Operation returning true. If fn
-// alters stored intervals' sort relationships, future tree operation behaviors
-// are undefined.
 func (t *btree) Do(fn Operation) bool {
 	if t.root == nil {
 		return false
@@ -745,7 +728,9 @@ func (n *node) mergeWithRightChild(i int, fast bool) {
 	}
 }
 
-// btree is a B-tree based interval tree.
+var _ Tree = (*btree)(nil)
+
+// btree is an interval tree based on an augmented BTree.
 //
 // Tree stores Instances in an ordered structure, allowing easy insertion,
 // removal, and iteration.
@@ -802,9 +787,6 @@ func (n *node) rangeEnd() Comparable {
 	return maxEnd
 }
 
-// AdjustRanges fixes range fields for all nodes in the tree. This must be
-// called before Get, Do or DoMatching* is used if fast insertion or deletion
-// has been performed.
 func (t *btree) AdjustRanges() {
 	if t.isEmpty() {
 		return
@@ -835,8 +817,6 @@ func (t *btree) newNode() (n *node) {
 	return
 }
 
-// Insert inserts the Interface e into the tree. Insertions may replace an
-// existing Interface which is equal to the Interface e.
 func (t *btree) Insert(e Interface, fast bool) (err error) {
 	// t.metrics("Insert")
 	if err = isValidInterface(e); err != nil {
@@ -870,8 +850,6 @@ func (t *btree) Insert(e Interface, fast bool) (err error) {
 	return
 }
 
-// Delete deletes the Interface e if it exists in the B-tree. The deleted
-// Interface is equal to the Interface e.
 func (t *btree) Delete(e Interface, fast bool) (err error) {
 	// t.metrics("Delete")
 	if err = isValidInterface(e); err != nil {
@@ -895,21 +873,47 @@ func (t *btree) delete(e Interface, typ toRemove, fast bool) Interface {
 	return out
 }
 
-// Len returns the number of Interfaces currently in the tree.
 func (t *btree) Len() int {
 	return t.length
 }
 
+type stackElem struct {
+	node  *node
+	index int
+}
+
 type btreeIterator struct {
+	stack []*stackElem
 }
 
-// Next iterates over the items stored in the tree, in-order.
 func (ti *btreeIterator) Next() (i Interface, ok bool) {
-	panic("TODO")
+	if len(ti.stack) == 0 {
+		return nil, false
+	}
+	elem := ti.stack[len(ti.stack)-1]
+	curItem := elem.node.items[elem.index]
+	elem.index++
+	if elem.index >= len(elem.node.items) {
+		ti.stack = ti.stack[:len(ti.stack)-1]
+	}
+	if len(elem.node.children) > 0 {
+		for r := elem.node.children[elem.index]; r != nil; r = r.children[0] {
+			ti.stack = append(ti.stack, &stackElem{r, 0})
+			if len(r.children) == 0 {
+				break
+			}
+		}
+	}
+	return curItem, true
 }
 
-// Iterator creates an iterator to iterate over all intervals stored in the
-// tree, in-order.
 func (t *btree) Iterator() TreeIterator {
-	return &btreeIterator{}
+	var ti btreeIterator
+	for n := t.root; n != nil; n = n.children[0] {
+		ti.stack = append(ti.stack, &stackElem{n, 0})
+		if len(n.children) == 0 {
+			break
+		}
+	}
+	return &ti
 }

pkg/util/interval/btree_based_interval_test.go

@@ -29,7 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
 
-var btreeMinDegree = flag.Int("btree_min_degree", 32, "B-Tree minimum degree")
+var btreeMinDegree = flag.Int("btree_min_degree", DefaultBTreeMinimumDegree, "B-Tree minimum degree")
 
 func init() {
 	seed := timeutil.Now().Unix()
@@ -342,8 +342,11 @@ func checkEqualIntervals(t *testing.T, actual, expected items) {
 			break
 		}
 	}
-	sort.Sort(expected)
-	if !reflect.DeepEqual(actual, expected) {
+	itemsLen := len(expected)
+	sortedExpected := make(items, itemsLen)
+	copy(sortedExpected, expected)
+	sort.Sort(sortedExpected)
+	if !reflect.DeepEqual(actual, sortedExpected) {
 		t.Errorf("expected intervals %v, got %v", expected, actual)
 	}
 }
@@ -372,6 +375,15 @@ func checkTraversal(t *testing.T, tree *btree, ivs items) {
 	checkEqualIntervals(t, all, ivs)
 }
 
+func checkIterator(t *testing.T, tree *btree, ivs items) {
+	var actual items
+	it := tree.Iterator()
+	for r, ok := it.Next(); ok; r, ok = it.Next() {
+		actual = append(actual, r)
+	}
+	checkEqualIntervals(t, actual, ivs)
+}
+
 func checkFastDelete(t *testing.T, tree *btree, ivs items, deleteCount int) {
 	for i, iv := range ivs[:deleteCount] {
 		if err := tree.Delete(iv, true); err != nil {
@@ -536,18 +548,19 @@ func TestSmallTree(t *testing.T) {
 		if err := tree.Insert(iv, false); err != nil {
 			t.Fatalf("insert error: %s", err)
 		}
-		checkWithLen(t, &tree, i+1)
+		checkWithLen(t, tree, i+1)
 	}
 
-	checkTraversal(t, &tree, ivs)
+	checkTraversal(t, tree, ivs)
+	checkIterator(t, tree, ivs)
 
 	// Delete
 	l := tree.Len()
 	for i, iv := range ivs {
 		if err := tree.Delete(iv, false); err != nil {
 			t.Fatalf("delete error: %s", err)
 		}
-		checkWithLen(t, &tree, l-i-1)
+		checkWithLen(t, tree, l-i-1)
 	}
 }
 
@@ -562,10 +575,11 @@ func TestSmallTreeWithFastOperations(t *testing.T) {
 		}
 	}
 	tree.AdjustRanges()
-	checkWithLen(t, &tree, len(ivs))
+	checkWithLen(t, tree, len(ivs))
 
-	checkTraversal(t, &tree, ivs)
-	checkFastDelete(t, &tree, ivs, tree.Len())
+	checkTraversal(t, tree, ivs)
+	checkIterator(t, tree, ivs)
+	checkFastDelete(t, tree, ivs, tree.Len())
 }
 
 func TestLargeTree(t *testing.T) {
@@ -584,6 +598,23 @@ func TestLargeTree(t *testing.T) {
 		}
 	}
 	tree.AdjustRanges()
-	checkWithLen(t, &tree, treeSize)
-	checkFastDelete(t, &tree, ivs, 10)
+	checkWithLen(t, tree, treeSize)
+	checkFastDelete(t, tree, ivs, 10)
+}
+
+func TestIterator(t *testing.T) {
+	var ivs items
+	const treeSize = 400
+	for i := uint32(0); i < treeSize; i++ {
+		iv := makeMultiByteInterval(i, i+1, i)
+		ivs = append(ivs, iv)
+	}
+	tree := newBTreeWithDegree(InclusiveOverlapper, 4)
+	for _, iv := range ivs {
+		if err := tree.Insert(iv, true); err != nil {
+			t.Fatalf("fast insert error: %s", err)
+		}
+	}
+	tree.AdjustRanges()
+	checkIterator(t, tree, ivs)
 }

pkg/util/interval/interval.go

@@ -2,6 +2,9 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+// Package interval provides two implementations for an interval tree. One is
+// based on an augmented Left-Leaning Red Black tree. The other is based on an
+// augmented BTree.
 package interval
 
 import (
@@ -34,6 +37,13 @@ type Range struct {
 	Start, End Comparable
 }
 
+var _ = Range{}.Equal
+
+// Equal returns whether the two ranges are equal.
+func (r Range) Equal(other Range) bool {
+	return r.Start.Equal(other.Start) && r.End.Equal(other.End)
+}
+
 // String implements the Stringer interface.
 func (r Range) String() string {
 	return fmt.Sprintf("{%x-%x}", r.Start, r.End)
@@ -145,19 +155,51 @@ func (c Comparable) Equal(o Comparable) bool {
 // should traverse no further.
 type Operation func(Interface) (done bool)
 
+// Tree is an interval tree. For all the functions which have a fast argment,
+// fast being true means a fast operation which does not adjust node ranges. If
+// fast is false, node ranges are adjusted.
 type Tree interface {
+	// AdjustRanges fixes range fields for all nodes in the tree. This must be
+	// called before Get, Do or DoMatching* is used if fast insertion or deletion
+	// has been performed.
 	AdjustRanges()
+	// Len returns the number of intervals stored in the Tree.
 	Len() int
+	// Get returns a slice of Interfaces that overlap r in the tree. The slice is
+	// sorted nondecreasingly by interval start.
 	Get(r Range) []Interface
+	// GetWithOverlapper returns a slice of Interfaces that overlap r in the tree
+	// using the provided overlapper function. The slice is sorted nondecreasingly
+	// by interval start.
 	GetWithOverlapper(r Range, overlapper Overlapper) []Interface
+	// Insert inserts the Interface e into the tree. Insertions may replace an
+	// existing Interface which is equal to the Interface e.
 	Insert(e Interface, fast bool) error
+	// Delete deletes the Interface e if it exists in the tree. The deleted
+	// Interface is equal to the Interface e.
 	Delete(e Interface, fast bool) error
+	// Do performs fn on all intervals stored in the tree. The traversal is done
+	// in the nondecreasing order of interval start. A boolean is returned
+	// indicating whether the traversal was interrupted by an Operation returning
+	// true. If fn alters stored intervals' sort relationships, future tree
+	// operation behaviors are undefined.
 	Do(fn Operation) bool
+	// DoMatching performs fn on all intervals stored in the tree that overlaps r.
+	// The traversal is done in the nondecreasing order of interval start. A
+	// boolean is returned indicating whether the traversal was interrupted by an
+	// Operation returning true. If fn alters stored intervals' sort
+	// relationships, future tree operation behaviors are undefined.
 	DoMatching(fn Operation, r Range) bool
+	// Iterator creates an iterator to iterate over all intervals stored in the
+	// tree, in-order.
 	Iterator() TreeIterator
 }
 
+// TreeIterator iterates over all intervals stored in the interval tree, in-order.
 type TreeIterator interface {
+	// Next returns the current interval stored in the interval tree	and moves
+	// the iterator to the next interval. The method returns false if no intervals
+	// remain in the interval tree.
 	Next() (Interface, bool)
 }