Skip to content

Commit

Permalink
use for loops in iterx
Browse files Browse the repository at this point in the history
  • Loading branch information
mcy committed Feb 28, 2025
1 parent 7102ee9 commit 09a2279
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 80 deletions.
27 changes: 11 additions & 16 deletions internal/ext/iterx/consume.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
9 changes: 5 additions & 4 deletions internal/ext/iterx/filtermap.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
}
}

Expand Down
67 changes: 23 additions & 44 deletions internal/ext/iterx/get.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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.
Expand Down
22 changes: 9 additions & 13 deletions internal/ext/iterx/iterx.go
Original file line number Diff line number Diff line change
Expand Up @@ -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--
}
}
}

Expand All @@ -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
})
}
}
}
5 changes: 2 additions & 3 deletions internal/ext/mapsx/collect.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

0 comments on commit 09a2279

Please sign in to comment.