Merge pull request #151 from azr/master

augmentedtree: add traverse func

Author: dustinhiatt-wf

augmentedtree/atree.go

@@ -116,6 +116,22 @@ type tree struct {
 	dummy                node
 }
 
+func (t *tree) Traverse(fn func(id Interval)) {
+	nodes := []*node{t.root}
+
+	for len(nodes) != 0 {
+		c := nodes[len(nodes)-1]
+		fn(c.interval)
+		nodes = nodes[:len(nodes)-1]
+		if c.children[0] != nil {
+			nodes = append(nodes, c.children[0])
+		}
+		if c.children[1] != nil {
+			nodes = append(nodes, c.children[1])
+		}
+	}
+}
+
 func (tree *tree) resetDummy() {
 	tree.dummy.children[0], tree.dummy.children[1] = nil, nil
 	tree.dummy.red = false

augmentedtree/atree_test.go

@@ -622,3 +622,20 @@ func TestInsertDuplicateIntervalChildren(t *testing.T) {
 	result := tree.Query(constructSingleDimensionInterval(0, 10, 0))
 	assert.Contains(t, result, iv1)
 }
+
+func TestTraverse(t *testing.T) {
+	tree := newTree(1)
+	top := 30
+	for i := 0; i <= top; i++ {
+		tree.Add(constructSingleDimensionInterval(int64(i*10), int64((i+1)*10), uint64(i)))
+	}
+	found := map[uint64]bool{}
+	tree.Traverse(func(id Interval) {
+		found[id.ID()] = true
+	})
+	for i := 0; i <= top; i++ {
+		if found, _ := found[uint64(i)]; !found {
+			t.Errorf("could not find expected interval %d", i)
+		}
+	}
+}

augmentedtree/interface.go

@@ -69,4 +69,7 @@ type Tree interface {
 	// interval.  The provided interval's ID method is ignored so the
 	// provided ID is irrelevant.
 	Query(interval Interval) Intervals
+	// Traverse will traverse tree and give alls intervals
+	// found in an undefined order
+	Traverse(func(Interval))
 }