fix(treeset,iter): use context to control iterate(#52)

go.mod

@@ -2,6 +2,9 @@ module github.com/hashicorp/go-set
 
 go 1.18
 
-require github.com/shoenig/test v0.6.3
+require (
+	github.com/shoenig/test v0.6.3
+	go.uber.org/goleak v1.2.1
+)
 
 require github.com/google/go-cmp v0.5.9 // indirect

go.sum

@@ -1,4 +1,10 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/shoenig/test v0.6.3 h1:GVXWJFk9PiOjN0KoJ7VrJGH6uLPnqxR7/fe3HUPfE0c=
 github.com/shoenig/test v0.6.3/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
+github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
+go.uber.org/goleak v1.2.1 h1:NBol2c7O1ZokfZ0LEU9K6Whx/KnwvepVetCUhtKja4A=
+go.uber.org/goleak v1.2.1/go.mod h1:qlT2yGI9QafXHhZZLxlSuNsMw3FFLxBr+tBRlmO1xH4=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

treeset.go

@@ -4,6 +4,7 @@
 package set
 
 import (
+	"context"
 	"fmt"
 )
 
@@ -304,8 +305,9 @@ func (s *TreeSet[T, C]) Empty() bool {
 // Slice returns the elements of s as a slice, in order.
 func (s *TreeSet[T, C]) Slice() []T {
 	result := make([]T, 0, s.Size())
-	s.infix(func(n *node[T]) {
+	s.infix(func(n *node[T]) bool {
 		result = append(result, n.element)
+		return true
 	}, s.root)
 	return result
 }
@@ -322,10 +324,13 @@ func (s *TreeSet[T, C]) Subset(o *TreeSet[T, C]) bool {
 	if s.Size() < o.Size() {
 		return false
 	}
+	ctx, cl := context.WithCancel(context.Background())
+	defer cl()
+
 	// iterate o, and increment s finding each element
 	// i.e. merge algorithm but with channels
-	iterO := o.iterate()
-	iterS := s.iterate()
+	iterO := o.iterate(ctx)
+	iterS := s.iterate(ctx)
 
 	idxO := 0
 	idxS := 0
@@ -410,8 +415,10 @@ func (s *TreeSet[T, C]) Equal(o *TreeSet[T, C]) bool {
 		return false
 	}
 
-	iterS := s.iterate()
-	iterO := o.iterate()
+	ctx, cl := context.WithCancel(context.Background())
+	defer cl()
+	iterS := s.iterate(ctx)
+	iterO := o.iterate(ctx)
 	for i := 0; i < s.Size(); i++ {
 		nextS := <-iterS
 		nextO := <-iterO
@@ -443,8 +450,9 @@ func (s *TreeSet[T, C]) String() string {
 // element into a string. The result contains elements in order.
 func (s *TreeSet[T, C]) StringFunc(f func(element T) string) string {
 	l := make([]string, 0, s.Size())
-	s.infix(func(n *node[T]) {
+	s.infix(func(n *node[T]) bool {
 		l = append(l, f(n.element))
+		return true
 	}, s.root)
 	return fmt.Sprintf("%s", l)
 }
@@ -849,12 +857,14 @@ func (s *TreeSet[T, C]) compare(a, b *node[T]) int {
 	return s.comparison(a.element, b.element)
 }
 
-func (s *TreeSet[T, C]) infix(visit func(*node[T]), n *node[T]) {
+func (s *TreeSet[T, C]) infix(visit func(*node[T]) (next bool), n *node[T]) {
 	if n == nil {
 		return
 	}
 	s.infix(visit, n.left)
-	visit(n)
+	if !visit(n) {
+		return
+	}
 	s.infix(visit, n.right)
 }
 
@@ -903,12 +913,23 @@ func (s *TreeSet[T, C]) prefix(visit func(*node[T]), n *node[T]) {
 	s.prefix(visit, n.right)
 }
 
-func (s *TreeSet[T, C]) iterate() <-chan *node[T] {
+func (s *TreeSet[T, C]) iterate(ctx context.Context) <-chan *node[T] {
 	c := make(chan *node[T], 1)
-	v := func(n *node[T]) {
-		c <- n
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	v := func(n *node[T]) bool {
+		select {
+		case <-ctx.Done():
+			return false
+		case c <- n:
+			return true
+		}
 	}
-	go s.infix(v, s.root)
+	go func() {
+		defer close(c)
+		s.infix(v, s.root)
+	}()
 	return c
 }
 

treeset_test.go

@@ -4,12 +4,13 @@
 package set
 
 import (
+	"context"
 	"fmt"
+	"github.com/shoenig/test/must"
+	"go.uber.org/goleak"
 	"math/rand"
 	"strings"
 	"testing"
-
-	"github.com/shoenig/test/must"
 )
 
 const (
@@ -269,6 +270,11 @@ func TestTreeSet_Subset(t *testing.T) {
 		t2 := TreeSetFrom[int, Compare[int]]([]int{5, 1, 2, 8, 3}, Cmp[int])
 		must.True(t, t1.Subset(t2))
 	})
+	t.Run("diff set", func(t *testing.T) {
+		t1 := TreeSetFrom[int, Compare[int]]([]int{1, 2, 3, 4, 5}, Cmp[int])
+		t2 := TreeSetFrom[int, Compare[int]]([]int{6, 7, 8, 9, 10}, Cmp[int])
+		must.False(t, t1.Subset(t2))
+	})
 }
 
 func TestTreeSet_Union(t *testing.T) {
@@ -844,3 +850,40 @@ func shuffle(s []int) []int {
 	}
 	return c
 }
+
+func TestTreeSet_infix(t *testing.T) {
+	ts := TreeSetFrom[int, Compare[int]]([]int{4, 7, 1, 5, 2, 8, 9, 3, 11, 13}, Cmp[int])
+	isOdd := func(n *node[int]) bool {
+		return n.element%2 == 1
+	}
+	odds := make([]int, 0, 5)
+	ts.infix(func(n *node[int]) bool {
+		if n.element > 8 {
+			return false
+		}
+		if isOdd(n) {
+			odds = append(odds, n.element)
+		}
+
+		return true
+	}, ts.root)
+	must.Eq(t, []int{1, 3, 5, 7}, odds)
+}
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func TestTreeSet_iterate(t *testing.T) {
+	s := TreeSetFrom[int, Compare[int]]([]int{4, 7, 1, 5, 2, 8, 9, 3, 11}, Cmp[int])
+	ctx, cl := context.WithCancel(context.Background())
+	defer cl()
+	ret := make([]int, 0, 9)
+	ch := s.iterate(ctx)
+	for n := range ch {
+		if n.element > 3 {
+			break
+		}
+		ret = append(ret, n.element)
+	}
+	must.Eq(t, []int{1, 2, 3}, ret)
+}