Add treap ascending Iterators.

treap_test.go

 package gtreap
 
 import (
         "bytes"
         "testing"
 )
 
 func stringCompare(a, b interface{}) int {
         return bytes.Compare([]byte(a.(string)), []byte(b.(string)))
 }
@@ skipping 78 matching lines @@
         }
 }
 
 func load(x *Treap, arr []string) *Treap {
         for i, s := range arr {
                 x = x.Upsert(s, i)
         }
         return x
 }
 
-func visitExpect(t *testing.T, x *Treap, start string, arr []string) {
+type treapVisitor func(x *Treap, pivot Item, visitor ItemVisitor)
+
+func visitExpect(t *testing.T, x *Treap, v treapVisitor, start string, arr []string) {
         n := 0
-        x.VisitAscend(start, func(i Item) bool {
+        v(x, start, func(i Item) bool {
                 if i.(string) != arr[n] {
                         t.Errorf("expected visit item: %v, saw: %v", arr[n], i)
                 }
                 n++
                 return true
         })
         if n != len(arr) {
                 t.Errorf("expected # visit callbacks: %v, saw: %v", len(arr), n)
+                panic(nil)
         }
 }
 
-func TestVisit(t *testing.T) {
+func testVisitImpl(t *testing.T, fn treapVisitor) {
         x := NewTreap(stringCompare)
-        visitExpect(t, x, "a", []string{})
+        visitExpect(t, x, fn, "a", []string{})
 
         x = load(x, []string{"e", "d", "c", "c", "a", "b", "a"})
 
         visitX := func() {
-                visitExpect(t, x, "a", []string{"a", "b", "c", "d", "e"})
-                visitExpect(t, x, "a1", []string{"b", "c", "d", "e"})
-                visitExpect(t, x, "b", []string{"b", "c", "d", "e"})
-                visitExpect(t, x, "b1", []string{"c", "d", "e"})
-                visitExpect(t, x, "c", []string{"c", "d", "e"})
-                visitExpect(t, x, "c1", []string{"d", "e"})
-                visitExpect(t, x, "d", []string{"d", "e"})
-                visitExpect(t, x, "d1", []string{"e"})
-                visitExpect(t, x, "e", []string{"e"})
-                visitExpect(t, x, "f", []string{})
+                visitExpect(t, x, fn, "a", []string{"a", "b", "c", "d", "e"})
+                visitExpect(t, x, fn, "a1", []string{"b", "c", "d", "e"})
+                visitExpect(t, x, fn, "b", []string{"b", "c", "d", "e"})
+                visitExpect(t, x, fn, "b1", []string{"c", "d", "e"})
+                visitExpect(t, x, fn, "c", []string{"c", "d", "e"})
+                visitExpect(t, x, fn, "c1", []string{"d", "e"})
+                visitExpect(t, x, fn, "d", []string{"d", "e"})
+                visitExpect(t, x, fn, "d1", []string{"e"})
+                visitExpect(t, x, fn, "e", []string{"e"})
+                visitExpect(t, x, fn, "f", []string{})
         }
         visitX()
 
         var y *Treap
         y = x.Upsert("f", 1)
         y = y.Delete("a")
         y = y.Upsert("cc", 2)
         y = y.Delete("c")
 
-        visitExpect(t, y, "a", []string{"b", "cc", "d", "e", "f"})
-        visitExpect(t, y, "a1", []string{"b", "cc", "d", "e", "f"})
-        visitExpect(t, y, "b", []string{"b", "cc", "d", "e", "f"})
-        visitExpect(t, y, "b1", []string{"cc", "d", "e", "f"})
-        visitExpect(t, y, "c", []string{"cc", "d", "e", "f"})
-        visitExpect(t, y, "c1", []string{"cc", "d", "e", "f"})
-        visitExpect(t, y, "d", []string{"d", "e", "f"})
-        visitExpect(t, y, "d1", []string{"e", "f"})
-        visitExpect(t, y, "e", []string{"e", "f"})
-        visitExpect(t, y, "f", []string{"f"})
-        visitExpect(t, y, "z", []string{})
+        visitExpect(t, y, fn, "a", []string{"b", "cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "a1", []string{"b", "cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "b", []string{"b", "cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "b1", []string{"cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "c", []string{"cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "c1", []string{"cc", "d", "e", "f"})
+        visitExpect(t, y, fn, "d", []string{"d", "e", "f"})
+        visitExpect(t, y, fn, "d1", []string{"e", "f"})
+        visitExpect(t, y, fn, "e", []string{"e", "f"})
+        visitExpect(t, y, fn, "f", []string{"f"})
+        visitExpect(t, y, fn, "z", []string{})
 
         // an uninitialized treap
         z := NewTreap(stringCompare)
 
         // a treap to force left traversal of min
         lmt := NewTreap(stringCompare)
         lmt = lmt.Upsert("b", 2)
         lmt = lmt.Upsert("a", 1)
 
         // The x treap should be unchanged.
@@ skipping 18 matching lines @@
                 t.Error("expected max of nil")
         }
         if lmt.Min() != "a" {
                 t.Errorf("expected min of a")
         }
         if lmt.Max() != "b" {
                 t.Errorf("expeced max of b")
         }
 }
 
-func visitExpectEndAtC(t *testing.T, x *Treap, start string, arr []string) {
+func TestVisit(t *testing.T) {
+        testVisitImpl(t, func(x *Treap, pivot Item, visitor ItemVisitor) {
+                x.VisitAscend(pivot, visitor)
+        })
+}
+
+func TestVisitIter(t *testing.T) {
+        testVisitImpl(t, func(x *Treap, pivot Item, visitor ItemVisitor) {
+                it := x.Iterator(pivot)
+                for {
+                        i, ok := it.Next()
+                        if !ok || !visitor(i) {
+                                return
+                        }
+                }
+        })
+}
+
+func visitExpectEndAtC(t *testing.T, x *Treap, fn treapVisitor, start string, arr []string) {
         n := 0
         x.VisitAscend(start, func(i Item) bool {
                 if stringCompare(i, "c") > 0 {
                         return false
                 }
                 if i.(string) != arr[n] {
                         t.Errorf("expected visit item: %v, saw: %v", arr[n], i)
                 }
                 n++
                 return true
         })
         if n != len(arr) {
                 t.Errorf("expected # visit callbacks: %v, saw: %v", len(arr), n)
         }
 }
 
-func TestVisitEndEarly(t *testing.T) {
+func testVisitEndEarlyImpl(t *testing.T, fn treapVisitor) {
         x := NewTreap(stringCompare)
-        visitExpectEndAtC(t, x, "a", []string{})
+        visitExpectEndAtC(t, x, fn, "a", []string{})
 
         x = load(x, []string{"e", "d", "c", "c", "a", "b", "a", "e"})
 
         visitX := func() {
-                visitExpectEndAtC(t, x, "a", []string{"a", "b", "c"})
-                visitExpectEndAtC(t, x, "a1", []string{"b", "c"})
-                visitExpectEndAtC(t, x, "b", []string{"b", "c"})
-                visitExpectEndAtC(t, x, "b1", []string{"c"})
-                visitExpectEndAtC(t, x, "c", []string{"c"})
-                visitExpectEndAtC(t, x, "c1", []string{})
-                visitExpectEndAtC(t, x, "d", []string{})
-                visitExpectEndAtC(t, x, "d1", []string{})
-                visitExpectEndAtC(t, x, "e", []string{})
-                visitExpectEndAtC(t, x, "f", []string{})
+                visitExpectEndAtC(t, x, fn, "a", []string{"a", "b", "c"})
+                visitExpectEndAtC(t, x, fn, "a1", []string{"b", "c"})
+                visitExpectEndAtC(t, x, fn, "b", []string{"b", "c"})
+                visitExpectEndAtC(t, x, fn, "b1", []string{"c"})
+                visitExpectEndAtC(t, x, fn, "c", []string{"c"})
+                visitExpectEndAtC(t, x, fn, "c1", []string{})
+                visitExpectEndAtC(t, x, fn, "d", []string{})
+                visitExpectEndAtC(t, x, fn, "d1", []string{})
+                visitExpectEndAtC(t, x, fn, "e", []string{})
+                visitExpectEndAtC(t, x, fn, "f", []string{})
         }
         visitX()
 }
 
+func TestVisitEndEarly(t *testing.T) {
+        testVisitEndEarlyImpl(t, func(x *Treap, pivot Item, visitor ItemVisitor) {
+                x.VisitAscend(pivot, visitor)
+        })
+}
+
+func TestVisitEndEarlyIter(t *testing.T) {
+        testVisitEndEarlyImpl(t, func(x *Treap, pivot Item, visitor ItemVisitor) {
+                it := x.Iterator(pivot)
+                for {
+                        i, ok := it.Next()
+                        if !ok || !visitor(i) {
+                                return
+                        }
+                }
+        })
+}
+
 func TestPriorityAfterUpsert(t *testing.T) {
         // See https://github.com/steveyen/gtreap/issues/3 found by icexin.
 
         var check func(n *node, level int, expectedPriority map[string]int)
         check = func(n *node, level int, expectedPriority map[string]int) {
                 if n == nil {
                         return
                 }
                 if n.priority != expectedPriority[n.item.(string)] {
                         t.Errorf("wrong priority")
@@ skipping 14 matching lines @@
         })
 
         s = s.Upsert("m", 4)
 
         check(s.root, 0, map[string]int{
                 "m": 20,
                 "l": 18,
                 "n": 19,
         })
 }