jackhammer

iterx.go (3353B)

---

 package iterx
 
 import (
 	"iter"
 	"runtime"
 	"sync"
 
 	"git.sr.ht/~jackmordaunt/jackhammer/channel"
 	"git.sr.ht/~jackmordaunt/jackhammer/errorsx"
 )
 
 // Fanout processes a sequence concurrently.
 func Fanout[V any](seq iter.Seq[V], fn func(value V)) {
 	wg := sync.WaitGroup{}
 	defer wg.Wait()
 
 	for v := range seq {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
 			fn(v)
 		}()
 	}
 }
 
 // Fanout2 processes a sequence concurrently.
 func Fanout2[K any, V any](seq iter.Seq2[K, V], fn func(key K, value V)) {
 	type pair[K any, V any] struct {
 		key   K
 		value V
 	}
 
 	queue := make(chan pair[K, V])
 	wg := sync.WaitGroup{}
 
 	for ii := 0; ii < runtime.NumCPU(); ii += 1 {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
 			for p := range queue {
 				fn(p.key, p.value)
 			}
 		}()
 	}
 
 	for k, v := range seq {
 		queue <- pair[K, V]{key: k, value: v}
 	}
 
 	close(queue)
 	wg.Wait()
 }
 
 // FanoutErr2 processes a sequence cuncurrently, returning an aggregate error.
 func FanoutErr2[K any, V any](seq iter.Seq2[K, V], fn func(key K, value V) error) error {
 	errs := []error{}
 	failures := make(chan error)
 	done := make(chan any)
 
 	Fanout2(seq, func(key K, value V) {
 		if err := fn(key, value); err != nil {
 			failures <- err
 		}
 	})
 
 	go func() {
 		errs = channel.Collect(failures)
 		close(done)
 	}()
 
 	close(failures)
 	<-done
 
 	if len(errs) > 0 {
 		return errorsx.PolyError{Errs: errs}
 	}
 
 	return nil
 }
 
 // For2 processes a sequence with the given function.
 func For2[K any, V any](seq iter.Seq2[K, V], fn func(key K, value V)) {
 	for k, v := range seq {
 		fn(k, v)
 	}
 }
 
 // ForErr2 processes a sequence with the given function.
 func ForErr2[K any, V any](seq iter.Seq2[K, V], fn func(key K, value V) error) error {
 	var errs []error
 
 	for k, v := range seq {
 		if err := fn(k, v); err != nil {
 			errs = append(errs, err)
 		}
 	}
 	if len(errs) > 0 {
 		return errorsx.PolyError{Errs: errs}
 	}
 
 	return nil
 }
 
 // Map transforms an iterator by applying a function to each element and returning the result.
 func Map[S iter.Seq[In], In any, Out any](s S, fn func(In) Out) iter.Seq[Out] {
 	return func(yield func(Out) bool) {
 		for v := range s {
 			if !yield(fn(v)) {
 				break
 			}
 		}
 	}
 }
 
 // MapIndex is like [Map] but provides the iteration index to the transform function.
 func MapIndex[S iter.Seq[In], In any, Out any](s S, fn func(int, In) Out) iter.Seq[Out] {
 	return func(yield func(Out) bool) {
 		ii := -1
 		for v := range s {
 			ii += 1
 			if !yield(fn(ii, v)) {
 				break
 			}
 		}
 	}
 }
 
 // Map2 transforms an iterator by applying a function to each element and returning the result.
 func Map2[S iter.Seq2[In1, In2], In1 any, In2 any, Out any](s S, fn func(In1, In2) Out) iter.Seq[Out] {
 	return func(yield func(Out) bool) {
 		for v1, v2 := range s {
 			if !yield(fn(v1, v2)) {
 				break
 			}
 		}
 	}
 }
 
 // ToMap collects a sequence into a map, using [keyFor] to extract the key from each element.
 func ToMap[S iter.Seq2[int, E], E any, K comparable](s S, keyFor func(int, E) K) map[K]E {
 	out := map[K]E{}
 	for ii, v := range s {
 		out[keyFor(ii, v)] = v
 	}
 	return out
 }
 
 // WithIndex adapts a single-value iterator to a two-value iterator that yeilds the index.
 func WithIndex[T any](s iter.Seq[T]) iter.Seq2[int, T] {
 	return iter.Seq2[int, T](func(yield func(int, T) bool) {
 		i := 0
 		for v := range s {
 			if !yield(i, v) {
 				break
 			}
 			i++
 		}
 	})
 }