pipeline.go

// Package pipeline is a staged streaming dataflow runtime: composable
// push-based operators (Map / Filter / FlatMap / Window / MapReduce) over typed
// channels, with bounded backpressure, plus a flow-based-programming (FBP)
// Component/Port/Network layer for declaratively-wired pipelines.
//
// It is the dataflow-operator-graph that Go channels + a worker pool do not
// provide ergonomically. Module path dev.helix.pipeline.
package pipeline

import (
	"context"
	"sync"
)

// Stream is a push-based source of T. Subscribe drives every element to sink
// until the source is exhausted, the context is cancelled, or sink returns an
// error. Subscribe blocks until the stream completes.
type Stream[T any] interface {
	Subscribe(ctx context.Context, sink func(T) error) error
}

// Operator transforms a Stream[T] into a Stream[U] — a composable stage.
type Operator[T any, U any] func(in Stream[T]) Stream[U]

// streamFunc adapts a function to a Stream.
type streamFunc[T any] func(ctx context.Context, sink func(T) error) error

func (f streamFunc[T]) Subscribe(ctx context.Context, sink func(T) error) error {
	return f(ctx, sink)
}

// FromSlice builds a bounded Stream from a slice.
func FromSlice[T any](items []T) Stream[T] {
	return streamFunc[T](func(ctx context.Context, sink func(T) error) error {
		for _, it := range items {
			if err := ctx.Err(); err != nil {
				return err
			}
			if err := sink(it); err != nil {
				return err
			}
		}
		return nil
	})
}

// FromChannel builds a Stream from a receive channel.
func FromChannel[T any](ch <-chan T) Stream[T] {
	return streamFunc[T](func(ctx context.Context, sink func(T) error) error {
		for {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case v, ok := <-ch:
				if !ok {
					return nil
				}
				if err := sink(v); err != nil {
					return err
				}
			}
		}
	})
}

// Map applies fn to each element. An fn error aborts the stream.
func Map[T any, U any](fn func(T) (U, error)) Operator[T, U] {
	return func(in Stream[T]) Stream[U] {
		return streamFunc[U](func(ctx context.Context, sink func(U) error) error {
			return in.Subscribe(ctx, func(t T) error {
				u, err := fn(t)
				if err != nil {
					return err
				}
				return sink(u)
			})
		})
	}
}

// Filter keeps elements for which pred is true.
func Filter[T any](pred func(T) bool) Operator[T, T] {
	return func(in Stream[T]) Stream[T] {
		return streamFunc[T](func(ctx context.Context, sink func(T) error) error {
			return in.Subscribe(ctx, func(t T) error {
				if pred(t) {
					return sink(t)
				}
				return nil
			})
		})
	}
}

// FlatMap expands each element into zero or more outputs.
func FlatMap[T any, U any](fn func(T) []U) Operator[T, U] {
	return func(in Stream[T]) Stream[U] {
		return streamFunc[U](func(ctx context.Context, sink func(U) error) error {
			return in.Subscribe(ctx, func(t T) error {
				for _, u := range fn(t) {
					if err := sink(u); err != nil {
						return err
					}
				}
				return nil
			})
		})
	}
}

// Window batches elements into fixed-size slices; a trailing partial window is
// emitted on completion.
func Window[T any](size int) Operator[T, []T] {
	if size < 1 {
		size = 1
	}
	return func(in Stream[T]) Stream[[]T] {
		return streamFunc[[]T](func(ctx context.Context, sink func([]T) error) error {
			buf := make([]T, 0, size)
			err := in.Subscribe(ctx, func(t T) error {
				buf = append(buf, t)
				if len(buf) == size {
					out := buf
					buf = make([]T, 0, size)
					return sink(out)
				}
				return nil
			})
			if err != nil {
				return err
			}
			if len(buf) > 0 {
				return sink(buf)
			}
			return nil
		})
	}
}

// Pipe2 composes two operators into one stage.
func Pipe2[A, B, C any](op1 Operator[A, B], op2 Operator[B, C]) Operator[A, C] {
	return func(in Stream[A]) Stream[C] {
		return op2(op1(in))
	}
}

// Collect drains a Stream into a slice (terminal). Useful for tests + small
// terminal sinks.
func Collect[T any](ctx context.Context, s Stream[T]) ([]T, error) {
	var out []T
	var mu sync.Mutex
	err := s.Subscribe(ctx, func(t T) error {
		mu.Lock()
		out = append(out, t)
		mu.Unlock()
		return nil
	})
	return out, err
}

// MapReduce partitions the stream by key, reduces each partition, and returns
// the per-key reduced values. It is terminal (drains the input stream).
func MapReduce[T any, K comparable, R any](
	ctx context.Context,
	in Stream[T],
	mapper func(T) (K, R),
	reducer func(acc R, next R) R,
) (map[K]R, error) {
	acc := make(map[K]R)
	seen := make(map[K]bool)
	err := in.Subscribe(ctx, func(t T) error {
		k, r := mapper(t)
		if !seen[k] {
			acc[k] = r
			seen[k] = true
		} else {
			acc[k] = reducer(acc[k], r)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return acc, nil
}