parallel_slice_transform.go

package gofu

import "sync"

// Map applies fn to each element concurrently. [ Immutable ] [ time: O(n/p); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).Map(strconv.Itoa)
func (ps SliceParallel[T]) Map[R any](fn func(T) R) SliceParallel[R] {
	if ps.shouldSequential() {
		return SliceParallel[R]{Slice: ps.Slice.Map(fn), Scale: ps.Scale}
	}
	n := ps.Scale(len(ps.Slice))
	result := make(Slice[R], len(ps.Slice))
	var wg sync.WaitGroup
	start := 0
	ps.Slice.ForEachChunk(n, func(_ int, chunk Slice[T]) {
		end := start + chunk.Len()
		wg.Add(1)
		go func(chunk Slice[T], seg Slice[R]) {
			defer wg.Done()
			chunk.mapNoAlloc(fn, seg)
		}(chunk, result[start:end])
		start = end
	})
	wg.Wait()
	return SliceParallel[R]{Slice: result, Scale: ps.Scale}
}

// FlatMap applies fn to each element concurrently, then flattens the results. [ Immutable ] [ time: O(n/p+m); space: O(m)]
//
// Example:
//
//	bigSlice.Parallel(4).FlatMap(strings.Fields)
func (ps SliceParallel[T]) FlatMap[R any](fn func(T) []R) SliceParallel[R] {
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, s Slice[T]) SliceParallel[R] {
		return SliceParallel[R]{Slice: s.FlatMap(fn), Scale: ps.Scale}
	})
	out := make(Slice[R], 0)
	for _, r := range results {
		out = append(out, r.Slice...)
	}
	return SliceParallel[R]{Slice: out, Scale: ps.Scale}
}

// MapFilter applies fn to each element concurrently and filters out results where fn returns true. [ Immutable ] [ time: O(n/p); space: O(k)]
//
// Example:
//
//	bigSlice.Parallel(4).MapFilter(func(v int) (string, bool) { if v%2 == 0 { return strconv.Itoa(v), false }; return "", true })
func (ps SliceParallel[T]) MapFilter[R any](fn func(T) (R, bool)) SliceParallel[R] {
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, s Slice[T]) SliceParallel[R] {
		return SliceParallel[R]{Slice: s.MapFilter(fn), Scale: ps.Scale}
	})
	out := make(Slice[R], 0)
	for _, r := range results {
		out = append(out, r.Slice...)
	}
	return SliceParallel[R]{Slice: out, Scale: ps.Scale}
}

// Filter keeps only matching elements concurrently. [ Immutable ] [ time: O(n/p); space: O(k)]
//
// Example:
//
//	bigSlice.Parallel(4).Filter(func(v int) bool { return v%2 == 0 })
func (ps SliceParallel[T]) Filter(fn func(T) bool) SliceParallel[T] {
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, s Slice[T]) SliceParallel[T] {
		return SliceParallel[T]{Slice: s.Filter(fn), Scale: ps.Scale}
	})
	out := make(Slice[T], 0)
	for _, r := range results {
		out = append(out, r.Slice...)
	}
	return SliceParallel[T]{Slice: out, Scale: ps.Scale}
}

// Partition splits the list by fn into two parallel lists. [ Immutable ] [ time: O(n); space: O(n)]
//
// This method runs sequentially rather than in parallel to preserve the
// relative order of elements within each result list. For an unordered
// parallel version that may be faster on large data, see PartitionUnstable.
//
// Example:
//
//	evens, odds := bigSlice.Parallel(4).Partition(func(v int) bool { return v%2 == 0 })
func (ps SliceParallel[T]) Partition(fn func(T) bool) (SliceParallel[T], SliceParallel[T]) {
	a, b := ps.Slice.Partition(fn)
	return SliceParallel[T]{Slice: a, Scale: ps.Scale}, SliceParallel[T]{Slice: b, Scale: ps.Scale}
}

// PartitionUnstable splits the list by fn into two parallel lists. [ Immutable ] [ time: O(n/p); space: O(n)]
//
// Unlike Partition, this method does not guarantee element order within
// each result list. This allows parallel chunked execution: each chunk
// is partitioned independently, then all pass-chunks and fail-chunks
// are concatenated. Elements within each chunk maintain relative order,
// but chunk order is determined by goroutine scheduling.
//
// Use PartitionUnstable over Partition when you only need to split data
// and order within the two groups does not matter. For order-preserving
// behavior, use Partition.
//
// Example:
//
//	evens, odds := bigSlice.Parallel(4).PartitionUnstable(func(v int) bool { return v%2 == 0 })
func (ps SliceParallel[T]) PartitionUnstable(fn func(T) bool) (SliceParallel[T], SliceParallel[T]) {
	if ps.shouldSequential() {
		a, b := ps.Slice.Partition(fn)
		return SliceParallel[T]{Slice: a, Scale: ps.Scale}, SliceParallel[T]{Slice: b, Scale: ps.Scale}
	}
	type pr struct{ pass, fail Slice[T] }
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, chunk Slice[T]) pr {
		pass, fail := chunk.Partition(fn)
		return pr{pass, fail}
	})
	var totalPass, totalFail int
	for _, r := range results {
		totalPass += r.pass.Len()
		totalFail += r.fail.Len()
	}
	pass := make(Slice[T], totalPass)
	fail := make(Slice[T], totalFail)
	pi, fi := 0, 0
	for _, r := range results {
		copy(pass[pi:], r.pass)
		pi += r.pass.Len()
		copy(fail[fi:], r.fail)
		fi += r.fail.Len()
	}
	return SliceParallel[T]{Slice: pass, Scale: ps.Scale},
		SliceParallel[T]{Slice: fail, Scale: ps.Scale}
}

// DistinctBy removes duplicates by key concurrently. [ Immutable ] [ time: O(n/p); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).DistinctBy(func(v string) string { return strings.ToLower(v) })
func (ps SliceParallel[T]) DistinctBy[K comparable](fn func(T) K) SliceParallel[T] {
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, s Slice[T]) SliceParallel[T] {
		seq := s
		seen := make(map[K]struct{}, seq.Len())
		local := make(Slice[T], 0, seq.Len())
		for _, v := range seq {
			k := fn(v)
			if _, ok := seen[k]; !ok {
				seen[k] = struct{}{}
				local = append(local, v)
			}
		}
		return SliceParallel[T]{Slice: local, Scale: ps.Scale}
	})
	seen := make(map[K]struct{})
	out := make(Slice[T], 0, len(ps.Slice))
	for _, r := range results {
		for _, v := range r.Slice {
			k := fn(v)
			if _, ok := seen[k]; !ok {
				seen[k] = struct{}{}
				out = append(out, v)
			}
		}
	}
	return SliceParallel[T]{Slice: out, Scale: ps.Scale}
}

// GroupBy groups elements by key concurrently. [ Immutable ] [ time: O(n/p); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).GroupBy(func(v string) byte { return v[0] })
func (ps SliceParallel[T]) GroupBy[K comparable](fn func(T) K) Map[K, SliceParallel[T]] {
	results := ps.MapChunks(ps.Scale(len(ps.Slice)), func(_ int, s Slice[T]) Map[K, SliceParallel[T]] {
		m := s.GroupBy(fn)
		out := make(Map[K, SliceParallel[T]], len(m))
		for k, v := range m {
			out[k] = SliceParallel[T]{Slice: v, Scale: ps.Scale}
		}
		return out
	})
	out := make(Map[K, SliceParallel[T]])
	for _, m := range results {
		for k, v := range m {
			if existing, ok := out[k]; ok {
				out[k] = SliceParallel[T]{
					Slice: append(existing.Slice, v.Slice...),
					Scale: ps.Scale,
				}
			} else {
				out[k] = v
			}
		}
	}
	return out
}

// ToMap converts elements to key-value pairs concurrently. [ Immutable ] [ time: O(n/p); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).ToMap(func(v int) (string, int) { return strconv.Itoa(v), v })
func (ps SliceParallel[T]) ToMap[K comparable, V any](fn func(T) (K, V)) Map[K, V] {
	return collectMapResults(ps.MapChunks(ps.Scale(len(ps.Slice)),
		func(_ int, s Slice[T]) Map[K, V] { return s.ToMap(fn) },
	))
}

// Reverse returns a new parallel list with elements in reverse order. [ Immutable ] [ time: O(n); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).Reverse()
func (ps SliceParallel[T]) Reverse() SliceParallel[T] {
	return SliceParallel[T]{Slice: ps.Slice.Reverse(), Scale: ps.Scale}
}

// Shuffle returns a new parallel list with elements randomly shuffled. [ Immutable ] [ time: O(n); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).Shuffle()
func (ps SliceParallel[T]) Shuffle() SliceParallel[T] {
	return SliceParallel[T]{Slice: ps.Slice.Shuffle(), Scale: ps.Scale}
}

// SortBy returns a new parallel list sorted by fn. [ Immutable ] [ time: O(n log n); space: O(n)]
//
// Example:
//
//	bigSlice.Parallel(4).SortBy(func(a, b int) int { return a - b })
func (ps SliceParallel[T]) SortBy(fn func(T, T) int) SliceParallel[T] {
	return SliceParallel[T]{Slice: ps.Slice.SortBy(fn), Scale: ps.Scale}
}

func collectMapResults[K comparable, V any](results []Map[K, V]) Map[K, V] {
	if len(results) == 0 {
		return Map[K, V]{}
	}
	result := results[0]
	for _, m := range results[1:] {
		for k, v := range m {
			result[k] = v
		}
	}
	return result
}