From 0bb8bc9dbd9fe03192417a58cd7cd5668149a4d9 Mon Sep 17 00:00:00 2001 From: Edward Muller Date: Sun, 23 Mar 2025 19:29:36 -0700 Subject: [PATCH] more helpers --- seq.go | 313 ++++++++++++++++++++++++++++++++++++++++++++++++-- seq_test.go | 321 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 626 insertions(+), 8 deletions(-) diff --git a/seq.go b/seq.go index 485ce78..f2f11be 100644 --- a/seq.go +++ b/seq.go @@ -2,8 +2,10 @@ package seq import ( "cmp" + "context" "iter" "sync/atomic" + "time" ) // With returns a sequence with the provided values. The values are iterated over lazily when the returned sequence is iterated @@ -42,6 +44,58 @@ func WithKV[K, V any](kv ...KV[K, V]) iter.Seq2[K, V] { } } +// FromChan returns a sequence that yields values from the provided channel. The sequence is iterated over lazily when the +// returned sequence is iterated over. The sequence will end when the channel is closed. +// +// This allows for collecting values from a channel into a slice or similar relatively easily: +// +// s := slices.Collect(FromChan(ch)) +// // instead of +// var s []T +// for v := range ch { +// s = append(s, v) +// } +func FromChan[T any](ch <-chan T) iter.Seq[T] { + return func(yield func(T) bool) { + for t := range ch { + if !yield(t) { + return + } + } + } +} + +// ToChan returns a channel that yields values from the provided sequence. The provided sequence is iterated over lazily when +// the returned channel is iterated over. The channel is closed when the sequence is exhausted. +func ToChan[T any](seq iter.Seq[T]) <-chan T { + ch := make(chan T) + go func() { + defer close(ch) + for t := range seq { + ch <- t + } + }() + return ch +} + +// ToChanCtx returns a channel that yields values from the provided sequence. The provided sequence is iterated over +// lazily when the returned channel is iterated over. The channel is closed when the sequence is exhausted or the +// context is canceled, whichever comes first. +func ToChanCtx[T any](ctx context.Context, seq iter.Seq[T]) <-chan T { + ch := make(chan T) + go func() { + defer close(ch) + for t := range seq { + select { + case ch <- t: + case <-ctx.Done(): + return + } + } + }() + return ch +} + // Map the values in the sequence to a new sequence of values by applying the function fn to each value. Function application // happens lazily when the returned sequence is iterated over. func Map[T, O any](seq iter.Seq[T], fn func(T) O) iter.Seq[O] { @@ -545,22 +599,22 @@ func ContainsKV[K, V comparable](seq iter.Seq2[K, V], key K, value V) bool { return false } -// ContainsFunc returns true if the function returns true for any value in the sequence. The sequence is iterated over when -// ContainsFunc is called. -func ContainsFunc[T any](seq iter.Seq[T], equal func(T) bool) bool { +// ContainsFunc returns true if the predicate function returns true for any value in the sequence. The sequence is +// iterated over when ContainsFunc is called. +func ContainsFunc[T any](seq iter.Seq[T], predicate func(T) bool) bool { for t := range seq { - if equal(t) { + if predicate(t) { return true } } return false } -// ContainsKVFunc returns true if the function returns true for any key-value pair in the sequence. The sequence is iterated over -// when ContainsKVFunc is called. -func ContainsKVFunc[K, V any](seq iter.Seq2[K, V], equal func(K, V) bool) bool { +// ContainsKVFunc returns true if the predicate function returns true for any key-value pair in the sequence. The +// sequence is iterated over when ContainsKVFunc is called. +func ContainsKVFunc[K, V any](seq iter.Seq2[K, V], predicate func(K, V) bool) bool { for k, v := range seq { - if equal(k, v) { + if predicate(k, v) { return true } } @@ -696,3 +750,246 @@ func IsSortedKV[K, V cmp.Ordered](seq iter.Seq2[K, V]) bool { } return true } + +// Coalesce returns the first non zero value in the sequence. The provided sequence is iterated over before Coalesce +// returns. If no non-zero value is found, the second return value is false. +func Coalesce[T comparable](seq iter.Seq[T]) (T, bool) { + var value T + var found bool + for t := range seq { + if t != value { + return t, true + } + } + return value, found +} + +// CoalesceKV returns the first key-value pair in the sequence whose value is non zero. The provided sequence is +// iterated over before CoalesceKV returns. If no non-zero value is found, the second return value is false. +func CoalesceKV[K, V comparable](seq iter.Seq2[K, V]) (KV[K, V], bool) { + var value V + var found bool + for k, v := range seq { + if v != value { + return KV[K, V]{K: k, V: v}, true + } + } + return KV[K, V]{}, found +} + +// Count returns the number of elements in the sequence. The sequence is iterated over before Count returns. +func Count[T any](seq iter.Seq[T]) int { + var count int + for range seq { + count++ + } + return count +} + +// CountKV returns the number of key-value pairs in the sequence. The sequence is iterated over before CountKV returns. +func CountKV[K, V any](seq iter.Seq2[K, V]) int { + var count int + for range seq { + count++ + } + return count +} + +// CountBy returns the number of elements in the sequence for which the function returns true. The sequence is iterated over +// before CountBy returns. +func CountBy[T any](seq iter.Seq[T], fn func(T) bool) int { + var count int + for t := range seq { + if fn(t) { + count++ + } + } + return count +} + +// CountKVBy returns the number of key-value pairs in the sequence for which the function returns true. The sequence is +// iterated over before CountKVBy returns. +func CountKVBy[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) int { + var count int + for k, v := range seq { + if fn(k, v) { + count++ + } + } + return count +} + +// CountValues returns a key-value sequence where the keys are the values in the original sequence and the values are +// the number of times that value appears in the original sequence. The returned key-value sequence is unordered. The +// provided sequence is iterated over before CountValues returns. +func CountValues[T comparable](seq iter.Seq[T]) iter.Seq2[T, int] { + m := make(map[T]int) + for t := range seq { + m[t]++ + } + return func(yield func(T, int) bool) { + for k, v := range m { + if !yield(k, v) { + return + } + } + } +} + +// Drop n elements from the starts of the sequence. The provided sequence is iterated over lazily when the returned +// sequence is iterated over. +func Drop[T any](seq iter.Seq[T], n int) iter.Seq[T] { + return func(yield func(T) bool) { + for i, t := range IterKV(seq, IntK[T]()) { + if i < n { + continue + } + if !yield(t) { + return + } + } + } +} + +// DropKV n key-value pairs from the starts of the sequence. The provided sequence is iterated over lazily when the returned +// sequence is iterated over. +func DropKV[K, V any](seq iter.Seq2[K, V], n int) iter.Seq2[K, V] { + i := -1 + return func(yield func(K, V) bool) { + for k, v := range seq { + i++ + if i < n { + continue + } + if !yield(k, v) { + return + } + } + } +} + +// DropBy returns a sequence with all elements for which the function returns true removed. The provided sequence is +// iterated over lazily when the returned sequence is iterated over. This is the opposite of Filter. +func DropBy[T any](seq iter.Seq[T], fn func(T) bool) iter.Seq[T] { + return Filter(seq, func(t T) bool { + return !fn(t) + }) +} + +// DropKVBy returns a sequence with all key-value pairs for which the function returns true removed. The provided sequence +// is iterated over lazily when the returned sequence is iterated over. This is the opposite of FilterKV. +func DropKVBy[K, V any](seq iter.Seq2[K, V], fn func(K, V) bool) iter.Seq2[K, V] { + return FilterKV(seq, func(k K, v V) bool { + return !fn(k, v) + }) +} + +// EveryUntil returns a sequence that yields the time every d duration until the provided time. The ticker will adjust +// the time interval or drop ticks to make up for slow iteratee. The duration d must be greater than zero; if not, +// the function will panic. Waits d long before yielding the first element. +func EveryUntil(d time.Duration, until time.Time) iter.Seq[time.Time] { + return func(yield func(time.Time) bool) { + for now := range time.Tick(d) { + if now.After(until) { + return + } + if !yield(now) { + return + } + if now.After(until) { + return + } + } + } +} + +// EveryN returns a sequence that yields the time every d duration n times. The ticker will adjust the time interval or +// drop ticks to make up for slow iteratee. The duration d must be greater than zero; if not, the function will panic. +// Waits d long before yielding the first element. +func EveryN(d time.Duration, times int) iter.Seq[time.Time] { + return func(yield func(time.Time) bool) { + if times == 0 { + return + } + for now := range time.Tick(d) { + if !yield(now) { + return + } + times-- + if times == 0 { + return + } + } + } +} + +// MapToKV maps the values in the sequence to a new sequence of key-value pairs by applying the function fn to each value. Function +// application happens lazily when the returned sequence is iterated over. +func MapToKV[T, K, V any](seq iter.Seq[T], fn func(T) (K, V)) iter.Seq2[K, V] { + return func(yield func(K, V) bool) { + for t := range seq { + k, v := fn(t) + if !yield(k, v) { + return + } + } + } +} + +// Find returns the index of the first occurrence of the value in the sequence, the "index" (0 based) of the value, and true. If +// the value is not found, the first return value is the length of the sequence, the second return value is false. The provided +// sequence is iterated over when Find is called. +func Find[T comparable](seq iter.Seq[T], value T) (int, bool) { + var i int + var t T + for i, t = range IterKV(seq, IntK[T]()) { + if t == value { + return i, true + } + } + return i + 1, false +} + +// FindBy returns the first value in the sequence for which the function returns true, the "index" (0) based) of the +// value, and true. If no value is found, the first return value is the zero value of the type, the second return value +// is the length of the sequence, and the third return value is false. The provided sequence is iterated over when FindBy is called. +func FindBy[T any](seq iter.Seq[T], fn func(T) bool) (T, int, bool) { + var i int + var t T + for i, t = range IterKV(seq, IntK[T]()) { + if fn(t) { + return t, i, true + } + } + var z T + return z, i + 1, false +} + +// FindByKey returns the value of the first key-value pair in the sequence for which the function returns true, the +// "index" (0 based) of the value, and true. If the key is not found, the first return value is the zero value of the +// value type, the second return value is the length of the sequence, and the third return value is false. The provided +// sequence is iterated over when FindByKey is called. +func FindByKey[K comparable, V any](seq iter.Seq2[K, V], key K) (V, int, bool) { + var i int + for k, v := range seq { + if k == key { + return v, i, true + } + i++ + } + var v V + return v, i, false +} + +// FindByValue is like FindByKey, but returns the key of the first key-value pair whose value is equal to the provided value. +func FindByValue[K comparable, V comparable](seq iter.Seq2[K, V], value V) (K, int, bool) { + var i int + for k, v := range seq { + if v == value { + return k, i, true + } + i++ + } + var k K + return k, i, false +} diff --git a/seq_test.go b/seq_test.go index 29d4f8a..d7a7f95 100644 --- a/seq_test.go +++ b/seq_test.go @@ -1,10 +1,12 @@ package seq import ( + "context" "fmt" "slices" "strconv" "strings" + "time" ) func ExampleWith() { @@ -33,6 +35,7 @@ func ExampleWithKV() { // b 2 // c 3 } + func ExampleMap() { i := With(1, 2, 3) @@ -703,6 +706,18 @@ func ExampleRepeat() { // hi } +func ExampleRepeatKV() { + i := RepeatKV(3, "a", 1) + for k, v := range i { + fmt.Println(k, v) + } + + // Output: + // a 1 + // a 1 + // a 1 +} + func ExampleReplace() { i := With(1, 2, 3, 4, 5) @@ -767,3 +782,309 @@ func ExampleIsSortedKV() { // false // false } + +func ExampleFromChan() { + values := make(chan int, 3) + go func() { + for i := range 10 { + values <- i + } + close(values) + }() + + vals := slices.Collect(FromChan(values)) + fmt.Println(vals) + + // Output: + // [0 1 2 3 4 5 6 7 8 9] +} + +func ExampleToChan() { + i := With(1, 2, 3, 4, 5) + ch := ToChan(i) + + for v := range ch { + fmt.Println(v) + } + + // Output: + // 1 + // 2 + // 3 + // 4 + // 5 +} + +func ExampleToChanCtx() { + i := With(1, 2, 3, 4, 5) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := ToChanCtx(ctx, i) + + for v := range ch { + fmt.Println(v) + if v == 3 { + cancel() + } + } + + // Output: + // 1 + // 2 + // 3 +} + +func ExampleCoalesce() { + i := With(0, 0, 4, 5) + + fmt.Println(Coalesce(i)) + + // Output: + // 4 true +} + +func ExampleCoalesceKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 0}, tKV{K: "b", V: 0}, tKV{K: "c", V: 4}, tKV{K: "d", V: 5}) + + fmt.Println(CoalesceKV(i)) + + // Output: + // {c 4} true +} + +func ExampleCount() { + i := With(1, 2, 3, 4) + + fmt.Println(Count(i)) + + // Output: + // 4 +} + +func ExampleCountKV() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + fmt.Println(CountKV(i)) + + // Output: + // 3 +} + +func ExampleCountBy() { + i := With(1, 2, 3, 4, 5, 6, 7, 8, 9, 10) + + fmt.Println(CountBy(i, func(v int) bool { + return v%2 == 0 + })) + + // Output: + // 5 +} + +func ExampleCountKVBy() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + fmt.Println(CountKVBy(i, func(k string, v int) bool { + return v%2 == 0 + })) + + // Output: + // 1 +} + +func ExampleCountValues() { + i := With(1, 1, 2, 2, 3, 3, 3, 4) + + for k, v := range CountValues(i) { + fmt.Printf("%d: %v\n", k, v) + } + + // Unordered output: + // 1: 2 + // 2: 2 + // 3: 3 + // 4: 1 +} + +func ExampleDrop() { + i := With(1, 2, 3, 4, 5) + + for v := range Drop(i, 2) { + fmt.Println(v) + } + + for v := range Drop(i, 0) { + fmt.Println(v) + } + + // doesn't print anything + for v := range Drop(i, 100) { + fmt.Println(v) + } + + // Output: + // 3 + // 4 + // 5 + // 1 + // 2 + // 3 + // 4 + // 5 +} + +func ExampleDropKV() { + 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 DropKV(i, 2) { + fmt.Println(k, v) + } + + for k, v := range DropKV(i, 0) { + fmt.Println(k, v) + } + + // doesn't print anything + for k, v := range DropKV(i, 100) { + fmt.Println(k, v) + } + + // Output: + // c 3 + // a 1 + // b 2 + // c 3 +} + +func ExampleDropBy() { + i := With(1, 2, 3, 4, 5, 6, 7, 8, 9) + + s := DropBy(i, func(v int) bool { + return v%2 == 0 + }) + + fmt.Println(slices.Collect(s)) + + // Output: + // [1 3 5 7 9] +} + +func ExampleDropKVBy() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + s := DropKVBy(i, func(k string, v int) bool { + return v%2 == 0 + }) + + for k, v := range s { + fmt.Println(k, v) + } + + // Output: + // a 1 + // c 3 +} + +func ExampleEveryUntil() { + var i int + for t := range EveryUntil(time.Millisecond, time.Now().Add(10*time.Millisecond)) { + _ = t // 2025-03-23 18:53:05.064589166 -0700 PDT m=+0.007687209 + i++ + } + + fmt.Println(i) + + // Output: + // 9 +} + +func ExampleEveryN() { + var i int + for t := range EveryN(time.Millisecond, 10) { + _ = t // 2025-03-23 18:53:05.064589166 -0700 PDT m=+0.007687209 + i++ + } + + fmt.Println(i) + + // Output: + // 10 +} + +func ExampleMapToKV() { + i := With(1, 2, 3) + + for k, v := range MapToKV(i, func(i int) (string, int) { + return string([]byte{byte(64 + i)}), i + }) { + fmt.Println(k, v) + } + + // Output: + // A 1 + // B 2 + // C 3 +} + +func ExampleFind() { + i := With(1, 2, 3, 4, 5) + + fmt.Println(Find(i, 3)) + + fmt.Println(Find(i, 6)) + + // Output: + // 2 true + // 5 false +} + +func ExampleFindBy() { + i := With(1, 2, 3, 4, 5) + + v, idx, ok := FindBy(i, func(v int) bool { + return v == 3 + }) + + fmt.Println(v, idx, ok) + + v, idx, ok = FindBy(i, func(v int) bool { + return v == 6 + }) + + fmt.Println(v, idx, ok) + + // Output: + // 3 2 true + // 0 5 false +} + +func ExampleFindByKey() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + fmt.Println(FindByKey(i, "b")) + + fmt.Println(FindByKey(i, "d")) + + // Output: + // 2 1 true + // 0 3 false +} + +func ExampleFindByValue() { + type tKV = KV[string, int] + i := WithKV(tKV{K: "a", V: 1}, tKV{K: "b", V: 2}, tKV{K: "c", V: 3}) + + fmt.Println(FindByValue(i, 2)) + + fmt.Println(FindByValue(i, 4)) + + // Output: + // b 1 true + // 3 false +}