use for loops in iterx

Author: mcy

internal/ext/iterx/consume.go

@@ -26,37 +26,32 @@ import (
 // If p is nil, it is treated as func(_ T) bool { return true }.
 func Count[T any](seq iter.Seq[T]) int {
 	var total int
-	seq(func(_ T) bool {
+	for range seq {
 		total++
-		return true
-	})
+	}
 	return total
 }
 
 // Join is like [strings.Join], but works on an iterator. Elements are
 // stringified as if by [fmt.Print].
 func Join[T any](seq iter.Seq[T], sep string) string {
 	var out strings.Builder
-	first := true
-	seq(func(v T) bool {
-		if !first {
+	for i, v := range Enumerate(seq) {
+		if i > 0 {
 			out.WriteString(sep)
 		}
-		first = false
-
 		fmt.Fprint(&out, v)
-		return true
-	})
+	}
 	return out.String()
 }
 
 // Every returns whether every element of an iterator satisfies the given
 // predicate. Returns true if seq yields no values.
 func Every[T any](seq iter.Seq[T], p func(T) bool) bool {
-	all := true
-	seq(func(v T) bool {
-		all = p(v)
-		return all
-	})
-	return all
+	for v := range seq {
+		if !p(v) {
+			return false
+		}
+	}
+	return true
 }

internal/ext/iterx/filtermap.go

@@ -34,10 +34,11 @@ func Filter[T any](seq iter.Seq[T], p func(T) bool) iter.Seq[T] {
 // FilterMap combines the operations of [Map] and [Filter].
 func FilterMap[T, U any](seq iter.Seq[T], f func(T) (U, bool)) iter.Seq[U] {
 	return func(yield func(U) bool) {
-		seq(func(v T) bool {
-			v2, ok := f(v)
-			return !ok || yield(v2)
-		})
+		for v := range seq {
+			if v2, ok := f(v); ok && !yield(v2) {
+				return
+			}
+		}
 	}
 }
 

internal/ext/iterx/get.go

@@ -20,28 +20,23 @@ import (
 
 // First retrieves the first element of an iterator.
 func First[T any](seq iter.Seq[T]) (v T, ok bool) {
-	seq(func(x T) bool {
-		v = x
-		ok = true
-		return false
-	})
-	return v, ok
+	for x := range seq {
+		return x, true
+	}
+	return v, false
 }
 
 // OnlyOne retrieves the only element of an iterator.
 func OnlyOne[T any](seq iter.Seq[T]) (v T, ok bool) {
-	var found T
-	seq(func(x T) bool {
-		if !ok {
-			found = x
+	for i, x := range Enumerate(seq) {
+		if i > 0 {
+			var z T
+			// Ensure we return the zero value if there is more
+			// than one element.
+			return z, false
 		}
-		ok = !ok
-		return ok
-	})
-	if ok {
-		// Ensure we return the zero value if there is more
-		// than one element.
-		v = found
+		v = x
+		ok = true
 	}
 	return v, ok
 }
@@ -51,43 +46,27 @@ func OnlyOne[T any](seq iter.Seq[T]) (v T, ok bool) {
 // Returns the value and the index at which it was found, or -1 if it wasn't
 // found.
 func Find[T any](seq iter.Seq[T], p func(T) bool) (int, T) {
-	var v T
-	var idx int
-	var found bool
-	seq(func(x T) bool {
+	for i, x := range Enumerate(seq) {
 		if p(x) {
-			v = x
-			found = true
-			return false
+			return i, x
 		}
-		idx++
-		return true
-	})
-	if !found {
-		idx = -1
 	}
-	return idx, v
+	var z T
+	return -1, z
 }
 
 // Find2 is like [Find] but for two-element iterators.
 func Find2[T, U any](seq iter.Seq2[T, U], p func(T, U) bool) (int, T, U) {
-	var v1 T
-	var v2 U
-	var idx int
-	var found bool
-	seq(func(x1 T, x2 U) bool {
+	var i int
+	for x1, x2 := range seq {
 		if p(x1, x2) {
-			v1, v2 = x1, x2
-			found = true
-			return false
+			return i, x1, x2
 		}
-		idx++
-		return true
-	})
-	if !found {
-		idx = -1
+		i++
 	}
-	return idx, v1, v2
+	var z1 T
+	var z2 U
+	return -1, z1, z2
 }
 
 // Index returns the index of the first element of seq that satisfies p.

internal/ext/iterx/iterx.go

@@ -23,13 +23,12 @@ import (
 // Take returns an iterator over the first n elements of a sequence.
 func Take[T any](seq iter.Seq[T], n int) iter.Seq[T] {
 	return func(yield func(T) bool) {
-		seq(func(v T) bool {
-			if n > 0 && yield(v) {
-				n--
-				return true
+		for v := range seq {
+			if n == 0 || !yield(v) {
+				return
 			}
-			return false
-		})
+			n--
+		}
 	}
 }
 
@@ -56,15 +55,12 @@ func Strings[T any](seq iter.Seq[T]) iter.Seq[string] {
 // Chain returns an iterator that calls a sequence of iterators in sequence.
 func Chain[T any](seqs ...iter.Seq[T]) iter.Seq[T] {
 	return func(yield func(T) bool) {
-		var done bool
 		for _, seq := range seqs {
-			if done {
-				return
+			for v := range seq {
+				if !yield(v) {
+					return
+				}
 			}
-			seq(func(v T) bool {
-				done = !yield(v)
-				return !done
-			})
 		}
 	}
 }

internal/ext/mapsx/collect.go

@@ -25,10 +25,9 @@ func CollectSet[K comparable](seq iter.Seq[K]) map[K]struct{} {
 // InsertKeys is like [maps.Insert], but it implicitly fills in each map value
 // with the zero value.
 func InsertKeys[M ~map[K]V, K comparable, V any](m M, seq iter.Seq[K]) M {
-	seq(func(k K) bool {
+	for k := range seq {
 		var zero V
 		m[k] = zero
-		return true
-	})
+	}
 	return m
 }