diff --git a/README.md b/README.md index 7726a8a..601743f 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Golang's "missing" iterator/sequence functions. * `With(...T) iter.Seq[T]`: Construct a sequence using the provided values * `FromChan(<-chan T) iter.Seq[T]`: Returns a sequence that produces values until the channel is closed +* `FromChanCtx(context.Context, <-chan T) iter.Seq[T]`: Like FromChan but also stops when the context is canceled * `Repeat(int, T) iter.Seq[T]`: Returns a sequence which repeats the value n times ### iter.Seq2[K,V] @@ -27,6 +28,8 @@ Golang's "missing" iterator/sequence functions. * `IterK(iter.Seq2[K,V]) iter.Seq[K]`: Converts an iter.Seq2[K,V] to an iter.Seq[K] (keys only) * `IterV(iter.Seq2[K,V]) iter.Seq[V]`: Converts an iter.Seq2[K,V] to an iter.Seq[V] (values only) * `MapToKV(iter.Seq[T], func(T) (K,V)) iter.Seq2[K,V]`: Maps values to key-value pairs +* `SwapKV(iter.Seq2[K,V]) iter.Seq2[V,K]`: Swaps the keys and values of each pair +* `Enumerate(iter.Seq[T]) iter.Seq2[int,T]`: Pairs each value with its 0-based index; the index restarts on each iteration ## Transformation Functions @@ -34,6 +37,11 @@ Golang's "missing" iterator/sequence functions. * `Map(iter.Seq[T], func(T) O) iter.Seq[O]`: Maps the items in the sequence to another type * `MapKV(iter.Seq2[K,V], func(K,V) (K1,V1)) iter.Seq2[K1,V1]`: Maps the key-value pairs to other types +* `FlatMap(iter.Seq[T], func(T) iter.Seq[O]) iter.Seq[O]`: Maps each value to a sequence and yields the elements of each in order +* `Scan(iter.Seq[T], O, func(O,T) O) iter.Seq[O]`: Like Reduce but lazily yields the accumulated value after each element +* `ScanKV(iter.Seq2[K,V], O, func(O,K,V) O) iter.Seq[O]`: Like ReduceKV but lazily yields the accumulated value after each pair +* `Tap(iter.Seq[T], func(T)) iter.Seq[T]`: Yields the same elements, calling the function on each as it passes through +* `TapKV(iter.Seq2[K,V], func(K,V)) iter.Seq2[K,V]`: Yields the same pairs, calling the function on each as it passes through ### Filtering @@ -45,6 +53,19 @@ Golang's "missing" iterator/sequence functions. * `Append(iter.Seq[T], ...T) iter.Seq[T]`: Returns a new sequence with additional items appended * `AppendKV(iter.Seq2[K,V], ...KV[K,V]) iter.Seq2[K,V]`: Returns a new sequence with additional key-value pairs appended +### Combining + +* `Concat(...iter.Seq[T]) iter.Seq[T]`: Yields the elements of each sequence in order +* `ConcatKV(...iter.Seq2[K,V]) iter.Seq2[K,V]`: Yields the key-value pairs of each sequence in order +* `Zip(iter.Seq[A], iter.Seq[B]) iter.Seq2[A,B]`: Pairs the elements of two sequences positionally, ending at the shorter one +* `Merge(iter.Seq[T], iter.Seq[T]) iter.Seq[T]`: Merges two sorted sequences into one sorted sequence +* `MergeFunc(iter.Seq[T], iter.Seq[T], func(T,T) int) iter.Seq[T]`: Like Merge but uses a comparison function + +### Cycling + +* `Cycle(iter.Seq[T]) iter.Seq[T]`: Repeats the sequence forever (empty input yields an empty sequence) +* `CycleKV(iter.Seq2[K,V]) iter.Seq2[K,V]`: Repeats the key-value sequence forever (empty input yields an empty sequence) + ### Replacement * `Replace(iter.Seq[T], old, new T) iter.Seq[T]`: Replace old values with new values @@ -56,11 +77,30 @@ Golang's "missing" iterator/sequence functions. * `CompactFunc(iter.Seq[T], func(T,T) bool) iter.Seq[T]`: Like Compact but uses a function to compare elements * `CompactKV(iter.Seq2[K,V]) iter.Seq2[K,V]`: Yields all key-value pairs that are not equal to the previous pair * `CompactKVFunc(iter.Seq2[K,V], func(KV[K,V], KV[K,V]) bool) iter.Seq2[K,V]`: Like CompactKV but uses a function to compare pairs +* `Unique(iter.Seq[T]) iter.Seq[T]`: Yields the first occurrence of each distinct value (removes duplicates anywhere, not just adjacent) +* `UniqueKV(iter.Seq2[K,V]) iter.Seq2[K,V]`: Yields the first occurrence of each distinct key-value pair ### Chunking * `Chunk(iter.Seq[T], int) iter.Seq[iter.Seq[T]]`: Chunk the sequence into chunks of specified size * `ChunkKV(iter.Seq2[K,V], int) iter.Seq[iter.Seq2[K,V]]`: Chunk key-value pairs into chunks of specified size +* `Windows(iter.Seq[T], int) iter.Seq[iter.Seq[T]]`: Overlapping windows of the specified size (sliding by one element) +* `WindowsKV(iter.Seq2[K,V], int) iter.Seq[iter.Seq2[K,V]]`: Overlapping windows of key-value pairs +* `Flatten(iter.Seq[iter.Seq[T]]) iter.Seq[T]`: Yields the elements of each inner sequence in order (the inverse of Chunk) +* `FlattenKV(iter.Seq[iter.Seq2[K,V]]) iter.Seq2[K,V]`: Yields the key-value pairs of each inner sequence in order (the inverse of ChunkKV) + +### Grouping + +* `GroupBy(iter.Seq[T], func(T) K) iter.Seq2[K,[]T]`: Groups values by key in first-seen order +* `Partition(iter.Seq[T], func(T) bool) (iter.Seq[T], iter.Seq[T])`: Splits into matching and non-matching sequences +* `PartitionKV(iter.Seq2[K,V], func(K,V) bool) (iter.Seq2[K,V], iter.Seq2[K,V])`: Splits key-value pairs into matching and non-matching sequences + +### Taking + +* `Take(iter.Seq[T], int) iter.Seq[T]`: Take the first n elements of the sequence +* `TakeKV(iter.Seq2[K,V], int) iter.Seq2[K,V]`: Take the first n key-value pairs of the sequence +* `TakeWhile(iter.Seq[T], func(T) bool) iter.Seq[T]`: Take leading elements while the function returns true +* `TakeKVWhile(iter.Seq2[K,V], func(K,V) bool) iter.Seq2[K,V]`: Take leading key-value pairs while the function returns true ### Dropping @@ -68,6 +108,8 @@ Golang's "missing" iterator/sequence functions. * `DropKV(iter.Seq2[K,V], int) iter.Seq2[K,V]`: Drop n key-value pairs from the start of the sequence * `DropBy(iter.Seq[T], func(T) bool) iter.Seq[T]`: Drop all elements for which the function returns true * `DropKVBy(iter.Seq2[K,V], func(K,V) bool) iter.Seq2[K,V]`: Drop all key-value pairs for which the function returns true +* `DropWhile(iter.Seq[T], func(T) bool) iter.Seq[T]`: Drop leading elements while the function returns true, then yield the rest +* `DropKVWhile(iter.Seq2[K,V], func(K,V) bool) iter.Seq2[K,V]`: Drop leading key-value pairs while the function returns true, then yield the rest ## Aggregation Functions @@ -85,6 +127,12 @@ Golang's "missing" iterator/sequence functions. * `Reduce(iter.Seq[T], O, func(O,T) O) O`: Reduce the sequence to a single value * `ReduceKV(iter.Seq2[K,V], O, func(O,K,V) O) O`: Reduce key-value pairs to a single value +### Numeric + +* `Sum(iter.Seq[T]) T`: Sum of the values (zero for an empty sequence); T is any integer or float type +* `Product(iter.Seq[T]) T`: Product of the values (one for an empty sequence); T is any integer or float type +* `Average(iter.Seq[T]) (float64, bool)`: Arithmetic mean of the values; false if the sequence is empty + ### Counting * `Count(iter.Seq[T]) int`: Returns the number of elements in the sequence @@ -116,6 +164,13 @@ Golang's "missing" iterator/sequence functions. * `ContainsFunc(iter.Seq[T], func(T) bool) bool`: Returns true if predicate returns true for any value * `ContainsKVFunc(iter.Seq2[K,V], func(K,V) bool) bool`: Returns true if predicate returns true for any key-value pair +### Predicates + +* `All(iter.Seq[T], func(T) bool) bool`: Returns true if the function returns true for every value (true for empty) +* `AllKV(iter.Seq2[K,V], func(K,V) bool) bool`: Returns true if the function returns true for every key-value pair (true for empty) +* `None(iter.Seq[T], func(T) bool) bool`: Returns true if the function returns false for every value (true for empty) +* `NoneKV(iter.Seq2[K,V], func(K,V) bool) bool`: Returns true if the function returns false for every key-value pair (true for empty) + ### Finding * `Find(iter.Seq[T], T) (int, bool)`: Returns the index of the first occurrence of the value @@ -124,6 +179,8 @@ Golang's "missing" iterator/sequence functions. * `FindByValue(iter.Seq2[K,V], V) (K, int, bool)`: Returns the key of the first key-value pair with the given value * `At(iter.Seq[T], int) (T, bool)`: Returns the value at the given 0-based index, or zero value and false if out of range * `AtKV(iter.Seq2[K,V], int) (K, V, bool)`: Returns the key and value at the given 0-based index, or zero values and false if out of range +* `Last(iter.Seq[T]) (T, bool)`: Returns the final value in the sequence, or zero value and false if empty +* `LastKV(iter.Seq2[K,V]) (K, V, bool)`: Returns the final key-value pair in the sequence, or zero values and false if empty ## Utility Functions @@ -141,3 +198,4 @@ Golang's "missing" iterator/sequence functions. ## Types * `KV[K,V]`: A struct that pairs a key and value together for use with key-value sequence functions +* `Number`: A constraint permitting any integer or floating point type, used by Sum, Product, and Average diff --git a/seq.go b/seq.go index 763025f..cc77755 100644 --- a/seq.go +++ b/seq.go @@ -884,7 +884,10 @@ func EveryUntil(d time.Duration, until time.Time) iter.Seq[time.Time] { if !yield(now) { return } - if now.After(until) { + // Re-check the clock after the yield returns: a slow iteratee may have consumed the + // remaining time, and ending here beats waiting out another tick to notice. Checking + // now again would be useless — it cannot have changed since the check above. + if time.Now().After(until) { return } } @@ -1027,3 +1030,624 @@ func FindByValue[K comparable, V comparable](seq iter.Seq2[K, V], value V) (K, i var k K return k, i, false } + +// Take returns a sequence of the first n elements of the sequence. If the sequence has fewer than n elements, the +// returned sequence yields all of them. If n is not positive, the returned sequence is empty. The provided sequence is +// iterated over lazily when the returned sequence is iterated over. +func Take[T any](seq iter.Seq[T], n int) iter.Seq[T] { + return func(yield func(T) bool) { + if n <= 0 { + return + } + i := 0 + for t := range seq { + if !yield(t) { + return + } + i++ + if i == n { + return + } + } + } +} + +// TakeKV returns a sequence of the first n key-value pairs of the sequence. If the sequence has fewer than n pairs, the +// returned sequence yields all of them. If n is not positive, the returned sequence is empty. The provided sequence is +// iterated over lazily when the returned sequence is iterated over. +func TakeKV[K, V any](seq iter.Seq2[K, V], n int) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + if n <= 0 { + return + } + i := 0 + for k, v := range seq { + if !yield(k, v) { + return + } + i++ + if i == n { + return + } + } + } +} + +// TakeWhile returns a sequence of the leading elements of the sequence for which the function returns true. The +// sequence ends before the first element for which the function returns false. The provided sequence is iterated over +// lazily when the returned sequence is iterated over. +func TakeWhile[T any](seq iter.Seq[T], fn func(T) bool) iter.Seq[T] { + return func(yield func(T) bool) { + for t := range seq { + if !fn(t) || !yield(t) { + return + } + } + } +} + +// TakeKVWhile returns a sequence of the leading key-value pairs of the sequence for which the function returns true. +// The sequence ends before the first pair for which the function returns false. The provided sequence is iterated over +// lazily when the returned sequence is iterated over. +func TakeKVWhile[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for k, v := range seq { + if !fn(k, v) || !yield(k, v) { + return + } + } + } +} + +// DropWhile returns a sequence that skips the leading elements of the sequence for which the function returns true and +// then yields every remaining element, starting with the first element for which the function returns false. Unlike +// [DropBy], the function is not applied after the first non-matching element. The provided sequence is iterated over +// lazily when the returned sequence is iterated over. +func DropWhile[T any](seq iter.Seq[T], fn func(T) bool) iter.Seq[T] { + return func(yield func(T) bool) { + dropping := true + for t := range seq { + if dropping && fn(t) { + continue + } + dropping = false + if !yield(t) { + return + } + } + } +} + +// DropKVWhile returns a sequence that skips the leading key-value pairs of the sequence for which the function returns +// true and then yields every remaining pair, starting with the first pair for which the function returns false. Unlike +// [DropKVBy], the function is not applied after the first non-matching pair. The provided sequence is iterated over +// lazily when the returned sequence is iterated over. +func DropKVWhile[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + dropping := true + for k, v := range seq { + if dropping && fn(k, v) { + continue + } + dropping = false + if !yield(k, v) { + return + } + } + } +} + +// Concat returns a sequence that yields the elements of each provided sequence in order. The provided sequences are +// iterated over lazily when the returned sequence is iterated over. +func Concat[T any](seqs ...iter.Seq[T]) iter.Seq[T] { + return func(yield func(T) bool) { + for _, seq := range seqs { + for t := range seq { + if !yield(t) { + return + } + } + } + } +} + +// ConcatKV returns a sequence that yields the key-value pairs of each provided sequence in order. The provided +// sequences are iterated over lazily when the returned sequence is iterated over. +func ConcatKV[K, V any](seqs ...iter.Seq2[K, V]) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for _, seq := range seqs { + for k, v := range seq { + if !yield(k, v) { + return + } + } + } + } +} + +// Zip returns a sequence that pairs the elements of a and b positionally, yielding the elements of a as keys and the +// elements of b as values. The sequence ends when either input sequence ends. The provided sequences are iterated over +// lazily when the returned sequence is iterated over. +func Zip[A, B any](a iter.Seq[A], b iter.Seq[B]) iter.Seq2[A, B] { + return func(yield func(A, B) bool) { + next, stop := iter.Pull(b) + defer stop() + for av := range a { + bv, ok := next() + if !ok { + return + } + if !yield(av, bv) { + return + } + } + } +} + +// Merge merges two sorted sequences into one sorted sequence. [cmp.Compare] is used to compare elements. If the input +// sequences are not sorted the output will not be sorted either, but it will still contain every element of both. The +// provided sequences are iterated over lazily when the returned sequence is iterated over. +func Merge[T cmp.Ordered](a, b iter.Seq[T]) iter.Seq[T] { + return MergeFunc(a, b, cmp.Compare) +} + +// MergeFunc is like [Merge] but uses the function to compare elements. When elements compare equal, elements from b are +// yielded before elements from a. The provided sequences are iterated over lazily when the returned sequence is +// iterated over. +func MergeFunc[T any](a, b iter.Seq[T], compare func(T, T) int) iter.Seq[T] { + return func(yield func(T) bool) { + next, stop := iter.Pull(b) + defer stop() + bv, bok := next() + for av := range a { + for bok && compare(bv, av) <= 0 { + if !yield(bv) { + return + } + bv, bok = next() + } + if !yield(av) { + return + } + } + for bok { + if !yield(bv) { + return + } + bv, bok = next() + } + } +} + +// Flatten returns a sequence that yields the elements of each inner sequence in order. It is the inverse of [Chunk]. +// The provided sequence is iterated over lazily when the returned sequence is iterated over. +func Flatten[T any](seq iter.Seq[iter.Seq[T]]) iter.Seq[T] { + return func(yield func(T) bool) { + for inner := range seq { + for t := range inner { + if !yield(t) { + return + } + } + } + } +} + +// FlattenKV returns a sequence that yields the key-value pairs of each inner sequence in order. It is the inverse of +// [ChunkKV]. The provided sequence is iterated over lazily when the returned sequence is iterated over. +func FlattenKV[K, V any](seq iter.Seq[iter.Seq2[K, V]]) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for inner := range seq { + for k, v := range inner { + if !yield(k, v) { + return + } + } + } + } +} + +// FlatMap maps each value in the sequence to a sequence with the function and yields the elements of each resulting +// sequence in order. Function application happens lazily when the returned sequence is iterated over. +func FlatMap[T, O any](seq iter.Seq[T], fn func(T) iter.Seq[O]) iter.Seq[O] { + return func(yield func(O) bool) { + for t := range seq { + for o := range fn(t) { + if !yield(o) { + return + } + } + } + } +} + +// Unique returns a sequence that yields the first occurrence of each distinct value in the sequence. Unlike [Compact], +// which only removes adjacent duplicates, Unique removes duplicates anywhere in the sequence; it needs memory +// proportional to the number of distinct values to do so. The provided sequence is iterated over lazily when the +// returned sequence is iterated over. +func Unique[T comparable](seq iter.Seq[T]) iter.Seq[T] { + return func(yield func(T) bool) { + seen := make(map[T]struct{}) + for t := range seq { + if _, ok := seen[t]; ok { + continue + } + seen[t] = struct{}{} + if !yield(t) { + return + } + } + } +} + +// UniqueKV returns a sequence that yields the first occurrence of each distinct key-value pair in the sequence. Unlike +// [CompactKV], which only removes adjacent duplicates, UniqueKV removes duplicates anywhere in the sequence; it needs +// memory proportional to the number of distinct pairs to do so. The provided sequence is iterated over lazily when the +// returned sequence is iterated over. +func UniqueKV[K, V comparable](seq iter.Seq2[K, V]) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + seen := make(map[KV[K, V]]struct{}) + for k, v := range seq { + kv := KV[K, V]{K: k, V: v} + if _, ok := seen[kv]; ok { + continue + } + seen[kv] = struct{}{} + if !yield(k, v) { + return + } + } + } +} + +// Partition returns two sequences: the first yields the elements for which the function returns true, the second +// yields the rest. Each returned sequence iterates over the provided sequence independently, so iterating both +// iterates the provided sequence twice. +func Partition[T any](seq iter.Seq[T], fn func(T) bool) (iter.Seq[T], iter.Seq[T]) { + return Filter(seq, fn), DropBy(seq, fn) +} + +// PartitionKV returns two sequences: the first yields the key-value pairs for which the function returns true, the +// second yields the rest. Each returned sequence iterates over the provided sequence independently, so iterating both +// iterates the provided sequence twice. +func PartitionKV[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) (iter.Seq2[K, V], iter.Seq2[K, V]) { + return FilterKV(seq, fn), DropKVBy(seq, fn) +} + +// GroupBy returns a key-value sequence where the keys are the results of applying keyFn to each value and the values +// are slices of the values that produced each key, in encounter order. Keys are yielded in first-seen order. The +// provided sequence is iterated over completely when the returned sequence is iterated over. +func GroupBy[K comparable, T any](seq iter.Seq[T], keyFn func(T) K) iter.Seq2[K, []T] { + return func(yield func(K, []T) bool) { + groups := make(map[K][]T) + var order []K + for t := range seq { + k := keyFn(t) + if _, ok := groups[k]; !ok { + order = append(order, k) + } + groups[k] = append(groups[k], t) + } + for _, k := range order { + if !yield(k, groups[k]) { + return + } + } + } +} + +// Windows returns a sequence of overlapping windows of size consecutive elements. Each window after the first drops +// the oldest element of the previous window and appends the next element of the sequence. If the sequence has fewer +// than size elements the returned sequence is empty. The size must be at least 1; if not, the function will panic. The +// provided sequence is iterated over lazily when the returned sequence is iterated over. +func Windows[T any](seq iter.Seq[T], size int) iter.Seq[iter.Seq[T]] { + if size < 1 { + panic("seq: Windows size must be at least 1") + } + return func(yield func(iter.Seq[T]) bool) { + window := make([]T, 0, size) + for t := range seq { + if len(window) == size { + copy(window, window[1:]) + window[size-1] = t + } else { + window = append(window, t) + } + if len(window) == size { + w := make([]T, size) + copy(w, window) + if !yield(With(w...)) { + return + } + } + } + } +} + +// WindowsKV is like [Windows] but for key-value pairs. If the sequence has fewer than size pairs the returned sequence +// is empty. The size must be at least 1; if not, the function will panic. The provided sequence is iterated over lazily +// when the returned sequence is iterated over. +func WindowsKV[K, V any](seq iter.Seq2[K, V], size int) iter.Seq[iter.Seq2[K, V]] { + if size < 1 { + panic("seq: WindowsKV size must be at least 1") + } + return func(yield func(iter.Seq2[K, V]) bool) { + window := make([]KV[K, V], 0, size) + for k, v := range seq { + if len(window) == size { + copy(window, window[1:]) + window[size-1] = KV[K, V]{K: k, V: v} + } else { + window = append(window, KV[K, V]{K: k, V: v}) + } + if len(window) == size { + w := make([]KV[K, V], size) + copy(w, window) + if !yield(WithKV(w...)) { + return + } + } + } + } +} + +// All returns true if the function returns true for every value in the sequence. All returns true for an empty +// sequence. The sequence is iterated over until the function returns false when All is called. +func All[T any](seq iter.Seq[T], fn func(T) bool) bool { + for t := range seq { + if !fn(t) { + return false + } + } + return true +} + +// AllKV returns true if the function returns true for every key-value pair in the sequence. AllKV returns true for an +// empty sequence. The sequence is iterated over until the function returns false when AllKV is called. +func AllKV[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) bool { + for k, v := range seq { + if !fn(k, v) { + return false + } + } + return true +} + +// None returns true if the function returns false for every value in the sequence. None returns true for an empty +// sequence. This is the opposite of [ContainsFunc]. The sequence is iterated over until the function returns true when +// None is called. +func None[T any](seq iter.Seq[T], fn func(T) bool) bool { + return !ContainsFunc(seq, fn) +} + +// NoneKV returns true if the function returns false for every key-value pair in the sequence. NoneKV returns true for +// an empty sequence. This is the opposite of [ContainsKVFunc]. The sequence is iterated over until the function returns +// true when NoneKV is called. +func NoneKV[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) bool { + return !ContainsKVFunc(seq, fn) +} + +// Number is the constraint used by the numeric aggregation functions [Sum], [Product], and [Average]. It permits any +// integer or floating point type. +type Number interface { + ~int | ~int8 | ~int16 | ~int32 | ~int64 | + ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr | + ~float32 | ~float64 +} + +// Sum returns the sum of the values in the sequence, or zero if the sequence is empty. The sequence is iterated over +// before Sum returns. +func Sum[T Number](seq iter.Seq[T]) T { + var sum T + for t := range seq { + sum += t + } + return sum +} + +// Product returns the product of the values in the sequence, or one if the sequence is empty. The sequence is iterated +// over before Product returns. +func Product[T Number](seq iter.Seq[T]) T { + product := T(1) + for t := range seq { + product *= t + } + return product +} + +// Average returns the arithmetic mean of the values in the sequence. If the sequence is empty, the second return value +// is false. The sequence is iterated over before Average returns. +func Average[T Number](seq iter.Seq[T]) (float64, bool) { + var sum float64 + var count int + for t := range seq { + sum += float64(t) + count++ + } + if count == 0 { + return 0, false + } + return sum / float64(count), true +} + +// Last returns the final value in the sequence. If the sequence is empty, the second return value is false. The +// sequence is iterated over completely before Last returns. +func Last[T any](seq iter.Seq[T]) (T, bool) { + var last T + var found bool + for t := range seq { + last = t + found = true + } + return last, found +} + +// LastKV returns the final key-value pair in the sequence. If the sequence is empty, the third return value is false. +// The sequence is iterated over completely before LastKV returns. +func LastKV[K, V any](seq iter.Seq2[K, V]) (K, V, bool) { + var lk K + var lv V + var found bool + for k, v := range seq { + lk = k + lv = v + found = true + } + return lk, lv, found +} + +// Scan is like [Reduce] but returns a sequence that yields the accumulated value after each element instead of only +// the final value. The initial value itself is not yielded, so the returned sequence has as many elements as the +// provided one. The provided sequence is iterated over lazily when the returned sequence is iterated over. +func Scan[T, O any](seq iter.Seq[T], initial O, fn func(agg O, t T) O) iter.Seq[O] { + return func(yield func(O) bool) { + agg := initial + for t := range seq { + agg = fn(agg, t) + if !yield(agg) { + return + } + } + } +} + +// ScanKV is like [ReduceKV] but returns a sequence that yields the accumulated value after each key-value pair instead +// of only the final value. The initial value itself is not yielded, so the returned sequence has as many elements as +// the provided one has pairs. The provided sequence is iterated over lazily when the returned sequence is iterated +// over. +func ScanKV[K, V, O any](seq iter.Seq2[K, V], initial O, fn func(agg O, k K, v V) O) iter.Seq[O] { + return func(yield func(O) bool) { + agg := initial + for k, v := range seq { + agg = fn(agg, k, v) + if !yield(agg) { + return + } + } + } +} + +// Cycle returns a sequence that yields the elements of the sequence repeatedly, restarting from the beginning each +// time the provided sequence is exhausted. The returned sequence is infinite unless the provided sequence is empty, so +// bound iteration with something like [Take] or a break. The provided sequence must be re-iterable; single-use +// sequences (like those from [FromChan]) will not restart. +func Cycle[T any](seq iter.Seq[T]) iter.Seq[T] { + return func(yield func(T) bool) { + for { + empty := true + for t := range seq { + empty = false + if !yield(t) { + return + } + } + if empty { + return + } + } + } +} + +// CycleKV returns a sequence that yields the key-value pairs of the sequence repeatedly, restarting from the beginning +// each time the provided sequence is exhausted. The returned sequence is infinite unless the provided sequence is +// empty, so bound iteration with something like [TakeKV] or a break. The provided sequence must be re-iterable; +// single-use sequences will not restart. +func CycleKV[K, V any](seq iter.Seq2[K, V]) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for { + empty := true + for k, v := range seq { + empty = false + if !yield(k, v) { + return + } + } + if empty { + return + } + } + } +} + +// SwapKV returns a sequence with the keys and values of each pair swapped: the values become the keys and the keys +// become the values. The provided sequence is iterated over lazily when the returned sequence is iterated over. +func SwapKV[K, V any](seq iter.Seq2[K, V]) iter.Seq2[V, K] { + return func(yield func(V, K) bool) { + for k, v := range seq { + if !yield(v, k) { + return + } + } + } +} + +// Tap returns a sequence that yields the same elements as the provided sequence, calling the function on each element +// as it passes through. Useful for debugging or other side effects in the middle of a pipeline. The function is +// applied lazily when the returned sequence is iterated over. +func Tap[T any](seq iter.Seq[T], fn func(T)) iter.Seq[T] { + return func(yield func(T) bool) { + for t := range seq { + fn(t) + if !yield(t) { + return + } + } + } +} + +// TapKV returns a sequence that yields the same key-value pairs as the provided sequence, calling the function on each +// pair as it passes through. Useful for debugging or other side effects in the middle of a pipeline. The function is +// applied lazily when the returned sequence is iterated over. +func TapKV[K, V any](seq iter.Seq2[K, V], fn func(K, V)) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for k, v := range seq { + fn(k, v) + if !yield(k, v) { + return + } + } + } +} + +// FromChanCtx is like [FromChan] but stops when the context is canceled, even if the channel is blocked. The sequence +// ends when the channel is closed or the context is canceled, whichever comes first. Cancellation takes priority: once +// the context is canceled no further values are yielded, even if the channel has values ready. +func FromChanCtx[T any](ctx context.Context, ch <-chan T) iter.Seq[T] { + return func(yield func(T) bool) { + for { + // An already-canceled context must win over a ready channel; a bare select chooses randomly when + // both cases are ready. + select { + case <-ctx.Done(): + return + default: + } + select { + case <-ctx.Done(): + return + case t, ok := <-ch: + if !ok { + return + } + if !yield(t) { + return + } + } + } + } +} + +// Enumerate returns a key-value sequence that pairs each value in the sequence with its 0-based index. Unlike +// combining [IterKV] with [IntK], the index restarts at 0 each time the returned sequence is iterated over. The +// provided sequence is iterated over lazily when the returned sequence is iterated over. +func Enumerate[T any](seq iter.Seq[T]) iter.Seq2[int, T] { + return func(yield func(int, T) bool) { + var i int + for t := range seq { + if !yield(i, t) { + return + } + i++ + } + } +} diff --git a/seq_test.go b/seq_test.go index 6546e5b..f0f8c63 100644 --- a/seq_test.go +++ b/seq_test.go @@ -1,8 +1,10 @@ package seq import ( + "cmp" "context" "fmt" + "iter" "slices" "strconv" "strings" @@ -991,11 +993,16 @@ func ExampleDropKVBy() { } func ExampleEveryUntil() { - for t := range EveryUntil(time.Millisecond, time.Now().Add(10*time.Millisecond)) { + // The deadline is generous so the example stays deterministic; exact timing behaviors (tick counts, slow + // iteratees) are asserted in the stresstest package on a testing/synctest fake clock. + for t := range EveryUntil(time.Millisecond, time.Now().Add(time.Minute)) { _ = t // t == 2025-03-23 18:53:05.064589166 -0700 PDT m=+0.007687209 + fmt.Println("tick") + break } - // No output validation since this relies on time it will be flaky as a test + // Output: + // tick } func ExampleEveryN() { @@ -1232,3 +1239,970 @@ func ExampleIsSortedKV_negative() { // Output: // true } + +func ExampleTake() { + i := With(1, 2, 3, 4, 5) + + fmt.Println(slices.Collect(Take(i, 3))) + fmt.Println(slices.Collect(Take(i, 10))) + fmt.Println(slices.Collect(Take(i, 0))) + + for v := range Take(i, 3) { + fmt.Println(v) + break // stopping early stops the underlying sequence too + } + + // Output: + // [1 2 3] + // [1 2 3 4 5] + // [] + // 1 +} + +func ExampleTakeKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + for k, v := range TakeKV(i, 2) { + fmt.Println(k, v) + } + fmt.Println(CountKV(TakeKV(i, 0))) + + for k, v := range TakeKV(i, 3) { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // b 2 + // 0 + // a 1 +} + +func ExampleTakeWhile() { + i := With(1, 2, 3, 4, 1) + + s := TakeWhile(i, func(v int) bool { + return v < 3 + }) + + fmt.Println(slices.Collect(s)) + + for v := range TakeWhile(i, func(v int) bool { return v < 10 }) { + fmt.Println(v) + break + } + + // Output: + // [1 2] + // 1 +} + +func ExampleTakeKVWhile() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + s := TakeKVWhile(i, func(k string, v int) bool { + return v < 3 + }) + + for k, v := range s { + fmt.Println(k, v) + } + + for k, v := range TakeKVWhile(i, func(string, int) bool { return true }) { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // b 2 + // a 1 +} + +func ExampleDropWhile() { + i := With(1, 2, 3, 4, 1) + + // unlike DropBy, the trailing 1 is kept: dropping stops at the first non-matching element + fmt.Println(slices.Collect(DropWhile(i, func(v int) bool { + return v < 3 + }))) + + for v := range DropWhile(i, func(v int) bool { return v < 3 }) { + fmt.Println(v) + break + } + + // Output: + // [3 4 1] + // 3 +} + +func ExampleDropKVWhile() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 1}) + + s := DropKVWhile(i, func(k string, v int) bool { + return v < 2 + }) + + for k, v := range s { + fmt.Println(k, v) + } + + for k, v := range s { + fmt.Println(k, v) + break + } + + // Output: + // b 2 + // c 1 + // b 2 +} + +func ExampleConcat() { + i := Concat(With(1, 2), With(3), With[int](), With(4, 5)) + + fmt.Println(slices.Collect(i)) + + for v := range i { + fmt.Println(v) + break + } + + // Output: + // [1 2 3 4 5] + // 1 +} + +func ExampleConcatKV() { + type tKV = KV[string, int] + i := ConcatKV( + WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}), + WithKV(tKV{K: "c", V: 3}), + ) + + for k, v := range i { + fmt.Println(k, v) + } + + for k, v := range i { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // b 2 + // c 3 + // a 1 +} + +func ExampleZip() { + letters := With("a", "b", "c") + numbers := With(1, 2) + + // the sequence ends when either input ends + for k, v := range Zip(letters, numbers) { + fmt.Println(k, v) + } + for k, v := range Zip(numbers, letters) { + fmt.Println(k, v) + } + + for k, v := range Zip(letters, numbers) { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // b 2 + // 1 a + // 2 b + // a 1 +} + +func ExampleMerge() { + a := With(1, 3, 5) + b := With(2, 4) + + fmt.Println(slices.Collect(Merge(a, b))) + + // stopping early stops both inputs, wherever the next element comes from + m := Merge(With(2, 4), With(1, 3, 5)) + fmt.Println(slices.Collect(Take(m, 1))) + fmt.Println(slices.Collect(Take(m, 2))) + fmt.Println(slices.Collect(Take(Merge(With(1, 2), With(3, 4, 5)), 4))) + + // Output: + // [1 2 3 4 5] + // [1] + // [1 2] + // [1 2 3 4] +} + +func ExampleMergeFunc() { + desc := func(a, b int) int { + return cmp.Compare(b, a) + } + + fmt.Println(slices.Collect(MergeFunc(With(5, 3, 1), With(4, 2), desc))) + + // Output: + // [5 4 3 2 1] +} + +func ExampleFlatten() { + chunks := Chunk(With(1, 2, 3, 4, 5), 2) + + fmt.Println(slices.Collect(Flatten(chunks))) + + for v := range Flatten(chunks) { + fmt.Println(v) + break + } + + // Output: + // [1 2 3 4 5] + // 1 +} + +func ExampleFlattenKV() { + type tKV = KV[string, int] + chunks := ChunkKV(WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}), 2) + + for k, v := range FlattenKV(chunks) { + fmt.Println(k, v) + } + + for k, v := range FlattenKV(chunks) { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // b 2 + // c 3 + // a 1 +} + +func ExampleFlatMap() { + lines := With("hello world", "foo bar") + + words := FlatMap(lines, func(s string) iter.Seq[string] { + return With(strings.Fields(s)...) + }) + + fmt.Println(slices.Collect(words)) + + for w := range words { + fmt.Println(w) + break + } + + // Output: + // [hello world foo bar] + // hello +} + +func ExampleUnique() { + i := With(1, 2, 1, 3, 2, 4) + + fmt.Println(slices.Collect(Unique(i))) + + for v := range Unique(i) { + fmt.Println(v) + break + } + + // Output: + // [1 2 3 4] + // 1 +} + +func ExampleUniqueKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "a", V: 1}, tKV{K: "a", V: 2}, tKV{K: "b", V: 1}) + + for k, v := range UniqueKV(i) { + fmt.Println(k, v) + } + + for k, v := range UniqueKV(i) { + fmt.Println(k, v) + break + } + + // Output: + // a 1 + // a 2 + // b 1 + // a 1 +} + +func ExamplePartition() { + evens, odds := Partition(With(1, 2, 3, 4, 5), func(v int) bool { + return v%2 == 0 + }) + + fmt.Println(slices.Collect(evens)) + fmt.Println(slices.Collect(odds)) + + // Output: + // [2 4] + // [1 3 5] +} + +func ExamplePartitionKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + big, small := PartitionKV(i, func(k string, v int) bool { + return v > 1 + }) + + for k, v := range big { + fmt.Println(k, v) + } + for k, v := range small { + fmt.Println(k, v) + } + + // Output: + // b 2 + // c 3 + // a 1 +} + +func ExampleGroupBy() { + words := With("apple", "avocado", "banana", "blueberry", "cherry") + + groups := GroupBy(words, func(s string) string { + return s[:1] + }) + + for k, group := range groups { + fmt.Println(k, group) + } + + for k, group := range groups { + fmt.Println(k, len(group)) + break + } + + // Output: + // a [apple avocado] + // b [banana blueberry] + // c [cherry] + // a 2 +} + +func ExampleWindows() { + for w := range Windows(With(1, 2, 3, 4), 2) { + fmt.Println(slices.Collect(w)) + } + + // a sequence shorter than the window size yields no windows + fmt.Println(Count(Windows(With(1, 2), 3))) + + for w := range Windows(With(1, 2, 3, 4), 2) { + fmt.Println(slices.Collect(w)) + break + } + + // Output: + // [1 2] + // [2 3] + // [3 4] + // 0 + // [1 2] +} + +func ExampleWindowsKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + for w := range WindowsKV(i, 2) { + fmt.Println(slices.Collect(IterK(w))) + } + + for w := range WindowsKV(i, 2) { + fmt.Println(slices.Collect(IterV(w))) + break + } + + // Output: + // [a b] + // [b c] + // [1 2] +} + +func ExampleAll() { + even := func(v int) bool { return v%2 == 0 } + + fmt.Println(All(With(2, 4, 6), even)) + fmt.Println(All(With(2, 3, 6), even)) + fmt.Println(All(With[int](), even)) + + // Output: + // true + // false + // true +} + +func ExampleAllKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + fmt.Println(AllKV(i, func(k string, v int) bool { return v > 0 })) + fmt.Println(AllKV(i, func(k string, v int) bool { return v > 1 })) + + // Output: + // true + // false +} + +func ExampleNone() { + even := func(v int) bool { return v%2 == 0 } + + fmt.Println(None(With(1, 3, 5), even)) + fmt.Println(None(With(1, 2, 3), even)) + + // Output: + // true + // false +} + +func ExampleNoneKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + fmt.Println(NoneKV(i, func(k string, v int) bool { return v > 2 })) + fmt.Println(NoneKV(i, func(k string, v int) bool { return v > 1 })) + + // Output: + // true + // false +} + +func ExampleSum() { + fmt.Println(Sum(With(1, 2, 3))) + fmt.Println(Sum(With(1.5, 2.5))) + fmt.Println(Sum(With[int]())) + + // Output: + // 6 + // 4 + // 0 +} + +func ExampleProduct() { + fmt.Println(Product(With(1, 2, 3, 4))) + fmt.Println(Product(With[int]())) + + // Output: + // 24 + // 1 +} + +func ExampleAverage() { + avg, ok := Average(With(1, 2, 3, 4)) + fmt.Println(avg, ok) + + avg, ok = Average(With[int]()) + fmt.Println(avg, ok) + + // Output: + // 2.5 true + // 0 false +} + +func ExampleLast() { + v, ok := Last(With(1, 2, 3)) + fmt.Println(v, ok) + + v, ok = Last(With[int]()) + fmt.Println(v, ok) + + // Output: + // 3 true + // 0 false +} + +func ExampleLastKV() { + type tKV = KV[string, int] + + k, v, ok := LastKV(WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2})) + fmt.Println(k, v, ok) + + k, v, ok = LastKV(WithKV[string, int]()) + fmt.Println(k, v, ok) + + // Output: + // b 2 true + // 0 false +} + +func ExampleScan() { + i := With(1, 2, 3, 4) + + sums := Scan(i, 0, func(agg, v int) int { + return agg + v + }) + + fmt.Println(slices.Collect(sums)) + + for v := range sums { + fmt.Println(v) + break + } + + // Output: + // [1 3 6 10] + // 1 +} + +func ExampleScanKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + sums := ScanKV(i, 0, func(agg int, k string, v int) int { + return agg + v + }) + + fmt.Println(slices.Collect(sums)) + + for v := range sums { + fmt.Println(v) + break + } + + // Output: + // [1 3 6] + // 1 +} + +func ExampleCycle() { + fmt.Println(slices.Collect(Take(Cycle(With(1, 2, 3)), 7))) + + // cycling an empty sequence ends immediately instead of spinning forever + fmt.Println(Count(Cycle(With[int]()))) + + // Output: + // [1 2 3 1 2 3 1] + // 0 +} + +func ExampleCycleKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range TakeKV(CycleKV(i), 3) { + fmt.Println(k, v) + } + + fmt.Println(CountKV(CycleKV(WithKV[string, int]()))) + + // Output: + // a 1 + // b 2 + // a 1 + // 0 +} + +func ExampleSwapKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range SwapKV(i) { + fmt.Println(k, v) + } + + for k, v := range SwapKV(i) { + fmt.Println(k, v) + break + } + + // Output: + // 1 a + // 2 b + // 1 a +} + +func ExampleTap() { + var sum int + s := Tap(With(1, 2, 3), func(v int) { + sum += v + }) + + fmt.Println(slices.Collect(s)) + fmt.Println(sum) + + for range s { + break + } + fmt.Println(sum) + + // Output: + // [1 2 3] + // 6 + // 7 +} + +func ExampleTapKV() { + type tKV = KV[string, int] + var keys []string + s := TapKV(WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}), func(k string, v int) { + keys = append(keys, k) + }) + + fmt.Println(CountKV(s)) + fmt.Println(keys) + + for range s { + break + } + fmt.Println(keys) + + // Output: + // 2 + // [a b] + // [a b a] +} + +func ExampleFromChanCtx() { + ctx, cancel := context.WithCancel(context.Background()) + + ch := make(chan int, 3) + ch <- 1 + ch <- 2 + ch <- 3 + close(ch) + fmt.Println(slices.Collect(FromChanCtx(ctx, ch))) + + ch2 := make(chan int, 2) + ch2 <- 1 + ch2 <- 2 + for v := range FromChanCtx(ctx, ch2) { + fmt.Println(v) + break + } + + // canceling the context ends the sequence even though ch3 never produces a value or closes + ch3 := make(chan int) + cancel() + fmt.Println(slices.Collect(FromChanCtx(ctx, ch3))) + + // Output: + // [1 2 3] + // 1 + // [] +} + +func ExampleEnumerate() { + i := Enumerate(With("a", "b", "c")) + + for idx, v := range i { + fmt.Println(idx, v) + } + + // the index restarts at 0 on each iteration of the sequence + for idx, v := range i { + fmt.Println(idx, v) + break + } + + // Output: + // 0 a + // 1 b + // 2 c + // 0 a +} + +func ExampleEveryN_stopEarly() { + fmt.Println(Count(EveryN(time.Millisecond, 0))) + + for range EveryN(time.Millisecond, 5) { + break + } + fmt.Println("stopped") + + // Output: + // 0 + // stopped +} + +func ExampleToChanCtx_canceled() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + for range ToChanCtx(ctx, With(1, 2, 3)) { + } + fmt.Println("closed") + + // Output: + // closed +} + +func ExampleFromChan_stopEarly() { + ch := make(chan int, 2) + ch <- 1 + ch <- 2 + close(ch) + + for v := range FromChan(ch) { + fmt.Println(v) + break + } + + // Output: + // 1 +} + +func ExampleEqualKV_mismatch() { + type tKV = KV[string, int] + + fmt.Println(EqualKV(WithKV(tKV{K: "a", V: 1}), WithKV(tKV{K: "a", V: 2}))) + fmt.Println(EqualKV(WithKV(tKV{K: "a", V: 1}), WithKV(tKV{K: "b", V: 1}))) + + // Output: + // false + // false +} + +func ExampleCoalesce_allZero() { + fmt.Println(Coalesce(With(0, 0, 0))) + + // Output: + // 0 false +} + +func ExampleCoalesceKV_allZero() { + type tKV = KV[string, int] + + fmt.Println(CoalesceKV(WithKV(tKV{K: "a", V: 0}))) + + // Output: + // { 0} false +} + +func ExampleMap_stopEarly() { + fmt.Println(slices.Collect(Take(Map(With(1, 2, 3), strconv.Itoa), 1))) + + // Output: + // [1] +} + +func ExampleMapKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range TakeKV(MapKV(i, func(k string, v int) (string, int) { return k, v * 10 }), 1) { + fmt.Println(k, v) + } + + // Output: + // a 10 +} + +func ExampleAppend_stopEarly() { + fmt.Println(slices.Collect(Take(Append(With(1, 2), 3), 1))) + fmt.Println(slices.Collect(Take(Append(With(1), 2, 3), 2))) + + // Output: + // [1] + // [1 2] +} + +func ExampleAppendKV_stopEarly() { + type tKV = KV[string, int] + + for k, v := range TakeKV(AppendKV(WithKV(tKV{K: "a", V: 1}), tKV{K: "b", V: 2}), 1) { + fmt.Println(k, v) + } + for k, v := range TakeKV(AppendKV(WithKV(tKV{K: "a", V: 1}), tKV{K: "b", V: 2}, tKV{K: "c", V: 3}), 2) { + fmt.Println(k, v) + } + + // Output: + // a 1 + // a 1 + // b 2 +} + +func ExampleFilterKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + for k, v := range TakeKV(FilterKV(i, func(string, int) bool { return true }), 1) { + fmt.Println(k, v) + } + + // Output: + // a 1 +} + +func ExampleIterKV_stopEarly() { + for k, v := range TakeKV(IterKV(With("a", "b"), IntK[string]()), 1) { + fmt.Println(k, v) + } + + // Output: + // 0 a +} + +func ExampleIterK_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + fmt.Println(slices.Collect(Take(IterK(i), 1))) + + // Output: + // [a] +} + +func ExampleIterV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + fmt.Println(slices.Collect(Take(IterV(i), 1))) + + // Output: + // [1] +} + +func ExampleCompact_stopEarly() { + fmt.Println(slices.Collect(Take(Compact(With(1, 1, 2)), 1))) + + // Output: + // [1] +} + +func ExampleCompactFunc_stopEarly() { + fmt.Println(slices.Collect(Take(CompactFunc(With(1, 1, 2), func(a, b int) bool { return a == b }), 1))) + + // Output: + // [1] +} + +func ExampleCompactKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range TakeKV(CompactKV(i), 1) { + fmt.Println(k, v) + } + + // Output: + // a 1 +} + +func ExampleCompactKVFunc_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range TakeKV(CompactKVFunc(i, func(a, b tKV) bool { return a == b }), 1) { + fmt.Println(k, v) + } + + // Output: + // a 1 +} + +func ExampleChunk_stopEarly() { + for c := range Take(Chunk(With(1, 2, 3, 4), 2), 1) { + fmt.Println(slices.Collect(c)) + } + + // Output: + // [1 2] +} + +func ExampleChunkKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}, tKV{K: "d", V: 4}) + + for c := range Take(ChunkKV(i, 2), 1) { + fmt.Println(slices.Collect(IterK(c))) + } + + // Output: + // [a b] +} + +func ExampleRepeat_stopEarly() { + fmt.Println(slices.Collect(Take(Repeat(5, "x"), 1))) + + // Output: + // [x] +} + +func ExampleRepeatKV_stopEarly() { + for k, v := range TakeKV(RepeatKV(5, "a", 1), 1) { + fmt.Println(k, v) + } + + // Output: + // a 1 +} + +func ExampleReplace_stopEarly() { + fmt.Println(slices.Collect(Take(Replace(With(1, 2, 3), 1, 9), 1))) + + // Output: + // [9] +} + +func ExampleReplaceKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}) + + for k, v := range TakeKV(ReplaceKV(i, tKV{K: "a", V: 1}, tKV{K: "z", V: 9}), 1) { + fmt.Println(k, v) + } + + // Output: + // z 9 +} + +func ExampleCountValues_stopEarly() { + for k, v := range TakeKV(CountValues(With("x", "x")), 1) { + fmt.Println(k, v) + } + + // Output: + // x 2 +} + +func ExampleDrop_stopEarly() { + fmt.Println(slices.Collect(Take(Drop(With(1, 2, 3), 1), 1))) + + // Output: + // [2] +} + +func ExampleDropKV_stopEarly() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + for k, v := range TakeKV(DropKV(i, 1), 1) { + fmt.Println(k, v) + } + + // Output: + // b 2 +} + +func ExampleMapToKV_stopEarly() { + for k, v := range TakeKV(MapToKV(With(1, 2, 3), func(i int) (int, int) { return i, i * i }), 1) { + fmt.Println(k, v) + } + + // Output: + // 1 1 +} diff --git a/stresstest/stress_test.go b/stresstest/stress_test.go index e0c22c8..cdf1fe1 100644 --- a/stresstest/stress_test.go +++ b/stresstest/stress_test.go @@ -10,6 +10,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/freeformz/seq" @@ -63,6 +64,56 @@ func TestEveryUntilPanicsOnNonPositiveDuration(t *testing.T) { mustPanic(t, "EveryUntil d=-1", func() { seq.EveryUntil(-time.Second, time.Now()) }) } +func TestEveryUntilTickCount(t *testing.T) { + // On the synctest fake clock ticks land exactly on every interval, so the count is exact: ticks at 10, 20, 30, + // and 40ms; the 50ms tick is past the deadline. + synctest.Test(t, func(t *testing.T) { + var ticks int + for range seq.EveryUntil(10*time.Millisecond, time.Now().Add(45*time.Millisecond)) { + ticks++ + } + if ticks != 4 { + t.Errorf("EveryUntil yielded %d ticks, want 4", ticks) + } + }) +} + +func TestEveryUntilSlowIterateeEndsWithoutExtraTick(t *testing.T) { + // Regression: EveryUntil used to re-check the pre-yield timestamp after the yield returned, so a slow iteratee + // that consumed the remaining time still waited out one more tick (40ms total here) before noticing the + // deadline had passed. It must end as soon as the yield returns (at 30ms: first tick at 10ms + 20ms sleep). + synctest.Test(t, func(t *testing.T) { + start := time.Now() + var ticks int + for range seq.EveryUntil(10*time.Millisecond, start.Add(15*time.Millisecond)) { + ticks++ + time.Sleep(20 * time.Millisecond) + } + if ticks != 1 { + t.Errorf("EveryUntil yielded %d ticks, want 1", ticks) + } + if elapsed := time.Since(start); elapsed != 30*time.Millisecond { + t.Errorf("EveryUntil ended after %v, want 30ms", elapsed) + } + }) +} + +func TestEveryNTiming(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + start := time.Now() + var ticks int + for range seq.EveryN(10*time.Millisecond, 3) { + ticks++ + } + if ticks != 3 { + t.Errorf("EveryN yielded %d ticks, want 3", ticks) + } + if elapsed := time.Since(start); elapsed != 30*time.Millisecond { + t.Errorf("EveryN ended after %v, want 30ms", elapsed) + } + }) +} + func TestEveryNPanicsOnNonPositiveDuration(t *testing.T) { mustPanic(t, "EveryN d=0", func() { seq.EveryN(0, 1) }) mustPanic(t, "EveryN d=-1", func() { seq.EveryN(-time.Second, 1) }) @@ -193,3 +244,56 @@ func TestToChanCtxCancelClosesChannel(t *testing.T) { } }) } + +func TestWindowsPanicsOnNonPositiveSize(t *testing.T) { + mustPanic(t, "Windows size 0", func() { seq.Windows(seq.With(1, 2, 3), 0) }) + mustPanic(t, "Windows size -1", func() { seq.Windows(seq.With(1, 2, 3), -1) }) +} + +func TestWindowsKVPanicsOnNonPositiveSize(t *testing.T) { + type kv = seq.KV[string, int] + mustPanic(t, "WindowsKV size 0", func() { seq.WindowsKV(seq.WithKV(kv{K: "a", V: 1}), 0) }) + mustPanic(t, "WindowsKV size -1", func() { seq.WindowsKV(seq.WithKV(kv{K: "a", V: 1}), -1) }) +} + +func TestCycleEmptySequenceTerminates(t *testing.T) { + // Cycle restarts its input forever; an empty input must end the sequence instead of spinning. + withTimeout(t, 5*time.Second, func() { + for range seq.Cycle(seq.With[int]()) { + } + }) + withTimeout(t, 5*time.Second, func() { + for range seq.CycleKV(seq.WithKV[string, int]()) { + } + }) +} + +func TestFromChanCtxCancelUnblocks(t *testing.T) { + // FromChanCtx must end when the context is canceled even if the channel never produces another value. Inside + // the synctest bubble a regression shows up as a deadlock panic instead of needing a real-time timeout. + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + ch := make(chan int) + go func() { + time.Sleep(10 * time.Millisecond) + cancel() + }() + for range seq.FromChanCtx(ctx, ch) { + t.Error("FromChanCtx yielded a value from an empty channel") + } + }) +} + +func TestFromChanCtxCanceledCtxWinsOverReadyChannel(t *testing.T) { + // Regression: with a bare two-case select, an already-canceled context raced a ready channel and values could + // still be yielded after cancellation. Cancellation must take priority. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + ch := make(chan int, 1) + ch <- 1 + for range 100 { + for range seq.FromChanCtx(ctx, ch) { + t.Fatal("FromChanCtx yielded a value after the context was already canceled") + } + } +}