Fix iterator semantics. Improve treap traversal.

Author: lthibault

README.md

@@ -165,7 +165,7 @@ func main() {
     // To do this, we use an iterator.  Contrary to treaps, iterators are stateful and
     // mutable!  As such, they are NOT thread-safe.  However, multiple concurrent
     // iterators can traverse the same treap safely.
-    for iterator := handle.Iter(root); iterator.Next(); {
+    for iterator := handle.Iter(root); iterator.Node != nil; iterator.Next(); {
         fmt.Printf("%s %s: %d\n", iterator.Key, iterator.Value, iterator.Weight)
     }
 }

comparator_test.go

@@ -9,6 +9,8 @@ import (
 )
 
 func TestMaxTreap(t *testing.T) {
+	t.Parallel()
+
 	comp := treap.MaxTreap(treap.IntComparator)
 
 	for _, tc := range []struct {
@@ -38,6 +40,8 @@ func TestMaxTreap(t *testing.T) {
 }
 
 func TestIntComparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -68,6 +72,8 @@ func TestIntComparator(t *testing.T) {
 }
 
 func TestInt8Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -98,6 +104,8 @@ func TestInt8Comparator(t *testing.T) {
 }
 
 func TestInt16Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -128,6 +136,8 @@ func TestInt16Comparator(t *testing.T) {
 }
 
 func TestInt32Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -158,6 +168,8 @@ func TestInt32Comparator(t *testing.T) {
 }
 
 func TestInt64Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -188,6 +200,8 @@ func TestInt64Comparator(t *testing.T) {
 }
 
 func TestStringComparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -221,6 +235,8 @@ func TestStringComparator(t *testing.T) {
 }
 
 func TestByteComparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -254,6 +270,7 @@ func TestByteComparator(t *testing.T) {
 }
 
 func TestTimeComparator(t *testing.T) {
+	t.Parallel()
 
 	t0 := time.Now()
 
@@ -284,6 +301,8 @@ func TestTimeComparator(t *testing.T) {
 }
 
 func TestUIntComparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -314,6 +333,8 @@ func TestUIntComparator(t *testing.T) {
 }
 
 func TestUInt8Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -344,6 +365,8 @@ func TestUInt8Comparator(t *testing.T) {
 }
 
 func TestUInt16Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -374,6 +397,8 @@ func TestUInt16Comparator(t *testing.T) {
 }
 
 func TestUInt32Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -404,6 +429,8 @@ func TestUInt32Comparator(t *testing.T) {
 }
 
 func TestUInt64Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -434,6 +461,8 @@ func TestUInt64Comparator(t *testing.T) {
 }
 
 func TestFloat32Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}
@@ -464,6 +493,8 @@ func TestFloat32Comparator(t *testing.T) {
 }
 
 func TestFloat64Comparator(t *testing.T) {
+	t.Parallel()
+
 	for _, tc := range []struct {
 		desc string
 		test []interface{}

handle.go

@@ -214,10 +214,9 @@ func (h Handle) Pop(n *Node) (interface{}, *Node) {
 
 // Iter walks the tree in key-order.
 func (h Handle) Iter(n *Node) *Iterator {
-	return &Iterator{
-		n:     n,
-		stack: make([]*Node, 0, 16),
-	}
+	it := &Iterator{stack: push(nil, n)}
+	it.Next()
+	return it
 }
 
 func (h Handle) leftRotation(n *Node) *Node {
@@ -251,39 +250,3 @@ func (h Handle) rightRotation(n *Node) *Node {
 		n.Right.Right,
 	)
 }
-
-// Iterator contains treap iteration state.  Its methods are NOT thread-safe, but
-// multiple concurrent iterators are supported.
-type Iterator struct {
-	*Node
-	n     *Node
-	stack []*Node
-}
-
-// Next item.
-func (it *Iterator) Next() (more bool) {
-	for more = it.More(); more; {
-		if it.n != nil {
-			it.stack = append(it.stack, it.n)
-			it.n = it.n.Left
-			continue
-		}
-
-		it.n, it.stack = pop(it.stack)
-		it.Node = it.n
-		it.n = it.n.Right
-
-		break
-	}
-
-	return
-}
-
-// More returns true if the iterator has not traversed the entire treap.
-func (it *Iterator) More() bool {
-	return len(it.stack) > 0 || it.n != nil
-}
-
-func pop(stack []*Node) (*Node, []*Node) {
-	return stack[len(stack)-1], stack[:len(stack)-1]
-}

example_test.go renamed to handle_test.go

@@ -6,7 +6,7 @@ import (
 	"github.com/lthibault/treap"
 )
 
-func ExampleUsage() {
+func ExampleHandle() {
 	// Treap operations are performed by a lightweight handle.  Usually, you'll create a
 	// single global handle and share it between goroutines.  Handle's methods are thread-
 	// safe.
@@ -96,8 +96,9 @@ func ExampleUsage() {
 	// iterators can traverse the same treap safely.
 	var i int
 	fmt.Println("\n[ binary search-tree traversal (notice keys are sorted alphabetically)... ]")
-	for iterator := handle.Iter(root); iterator.Next(); i++ {
+	for iterator := handle.Iter(root); iterator.Node != nil; iterator.Next() {
 		fmt.Printf("[%d] %s %s: %d\n", i, iterator.Key, iterator.Value, iterator.Weight)
+		i++
 	}
 
 	// Output:

iterator.go

@@ -0,0 +1,53 @@
+package treap
+
+// Iterator contains treap iteration state.  Its methods are NOT thread-safe, but
+// multiple concurrent iterators are supported.
+type Iterator struct {
+	*Node
+	stack *stack
+}
+
+// Next item.
+func (it *Iterator) Next() {
+	// are we resuming?
+	if it.Node != nil {
+		it.Node = it.Node.Right
+	} else {
+		it.Node, it.stack = pop(it.stack)
+	}
+
+	for {
+		if it.Node == nil {
+			it.Node, it.stack = pop(it.stack)
+			break
+		}
+
+		it.stack = push(it.stack, it.Node)
+		it.Node = it.Node.Left
+	}
+}
+
+// Stack is a singly-linked list of nodes.
+type stack struct {
+	*Node
+	next *stack
+}
+
+func push(s *stack, n *Node) *stack {
+	if n == nil {
+		return s
+	}
+
+	return &stack{
+		Node: n,
+		next: s,
+	}
+}
+
+func pop(s *stack) (*Node, *stack) {
+	if s == nil {
+		return nil, nil
+	}
+
+	return s.Node, s.next
+}

iterator_test.go

@@ -0,0 +1,104 @@
+package treap_test
+
+import (
+	"testing"
+
+	"github.com/lthibault/treap"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIter_Empty(t *testing.T) {
+	t.Parallel()
+
+	var empty *treap.Node
+	assert.Nil(t, handle.Iter(empty).Node,
+		"iterator for empty root should have nil node")
+
+	var i int
+	for it := handle.Iter(empty); it.Node != nil; it.Next() {
+		i++
+	}
+	assert.Zero(t, i, "empty iterator shoudl not loop")
+}
+
+func TestIter_SingleEntry(t *testing.T) {
+	t.Parallel()
+
+	var root = &treap.Node{
+		Key:    1,
+		Value:  1,
+		Weight: 1,
+	}
+
+	assert.NotNil(t, handle.Iter(root).Node,
+		"iterator for non-empty root should have non-nil node")
+
+	var i int
+	for it := handle.Iter(root); it.Node != nil; it.Next() {
+		i++
+	}
+	assert.Equal(t, 1, i, "iterator should loop one time")
+}
+
+func TestIter_MultiEntry(t *testing.T) {
+	t.Parallel()
+
+	var (
+		root *treap.Node
+		ok   bool
+		tt   = []treap.Node{
+			{
+				Key:    0,
+				Value:  0,
+				Weight: 0,
+			},
+			{
+				Key:    1,
+				Value:  1,
+				Weight: 1,
+			},
+			{
+				Key:    2,
+				Value:  2,
+				Weight: 2,
+			},
+			{
+				Key:    3,
+				Value:  3,
+				Weight: 2,
+			},
+			{
+				Key:    4,
+				Value:  4,
+				Weight: 2,
+			},
+			{
+				Key:    5,
+				Value:  5,
+				Weight: 1,
+			},
+			{
+				Key:    6,
+				Value:  6,
+				Weight: 2,
+			},
+		}
+	)
+
+	for _, n := range tt {
+		root, ok = handle.Insert(root, n.Key, n.Value, n.Weight)
+		require.True(t, ok, "precondition failed: insert must succeed")
+	}
+
+	var ns []treap.Node
+	for it := handle.Iter(root); it.Node != nil; it.Next() {
+		t.Log(it.Node.Key)
+		ns = append(ns, *it.Node)
+	}
+	assert.Len(t, ns, len(tt), "iterator should loop %d times", len(tt))
+
+	for i, n := range tt {
+		assert.Equal(t, n.Key, ns[i].Key, "iterator should traverse in key order")
+	}
+}

treap_bench_test.go

@@ -166,7 +166,7 @@ func BenchmarkIterSync(b *testing.B) {
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		for it := handle.Iter(root); it.Next(); {
+		for it := handle.Iter(root); it.Node != nil; it.Next() {
 			discardNode = it.Node
 		}
 

treap_race_test.go

@@ -1,3 +1,4 @@
+//go:build race
 // +build race
 
 package treap_test