fix(treeset): infix can't stop correctly

zonewave · shoenig · commit b5d8a4651641 · 2023-05-13T08:08:01.000-05:00
diff --git a/treeset.go b/treeset.go
@@ -342,10 +342,10 @@ 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]) bool {
-		result = append(result, n.element)
+	s.ForEach(func(element T) bool {
+		result = append(result, element)
 		return true
-	}, s.root)
+	})
 	return result
 }
 
@@ -496,10 +496,10 @@ 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]) bool {
-		l = append(l, f(n.element))
+	s.ForEach(func(element T) bool {
+		l = append(l, f(element))
 		return true
-	}, s.root)
+	})
 	return fmt.Sprintf("%s", l)
 }
 
@@ -912,15 +912,17 @@ func (s *TreeSet[T, C]) compare(a, b *node[T]) int {
 // TreeNodeVisit is a function that is called for each node in the tree.
 type TreeNodeVisit[T any] func(*node[T]) (next bool)
 
-func (s *TreeSet[T, C]) infix(visit TreeNodeVisit[T], n *node[T]) {
+func (s *TreeSet[T, C]) infix(visit TreeNodeVisit[T], n *node[T]) (next bool) {
 	if n == nil {
+		return true
+	}
+	if next = s.infix(visit, n.left); !next {
 		return
 	}
-	s.infix(visit, n.left)
-	if !visit(n) {
+	if next = visit(n); !next {
 		return
 	}
-	s.infix(visit, n.right)
+	return s.infix(visit, n.right)
 }
 
 func (s *TreeSet[T, C]) fillLeft(n *node[T], k *[]T) {
diff --git a/treeset_test.go b/treeset_test.go
@@ -909,7 +909,7 @@ func TestTreeSet_infix(t *testing.T) {
 	}
 	odds := make([]int, 0, 5)
 	ts.infix(func(n *node[int]) bool {
-		if n.element > 8 {
+		if n.element == 8 {
 			return false
 		}
 		if isOdd(n) {

Original file line number	Diff line number	Diff line change
`@@ -909,7 +909,7 @@ func TestTreeSet_infix(t *testing.T) {`
`909`	`909`	`}`
`910`	`910`	`odds := make([]int, 0, 5)`
`911`	`911`	`ts.infix(func(n *node[int]) bool {`
`912`		`- if n.element > 8 {`
	`912`	`+ if n.element == 8 {`
`913`	`913`	`return false`
`914`	`914`	`}`
`915`	`915`	`if isOdd(n) {`