1 files changed, 81 insertions, 0 deletions
diff --git a/llmc/pool.go b/llmc/pool.go
new file mode 100644
index 0000000..c73d2d1
--- /dev/null
+++ b/llmc/pool.go
@@ -0,0 +1,81 @@
+package llmc
+
+import (
+	"iter"
+	"sync"
+)
+
+func Pool[In, Out any](seq iter.Seq[In], nWorkers int, f func(In) Out) <-chan Out {
+	out := make(chan Out, 256)
+
+	go func() {
+		defer close(out)
+
+		type work struct {
+			idx int
+			in  In
+		}
+
+		type result struct {
+			idx int
+			out Out
+		}
+
+		jobs := make(chan work, nWorkers)
+		results := make(chan result, nWorkers)
+
+		var wg sync.WaitGroup
+
+		for range nWorkers {
+			wg.Add(1)
+
+			go func() {
+				defer wg.Done()
+
+				for j := range jobs {
+					results <- result{j.idx, f(j.in)}
+				}
+			}()
+		}
+
+		go func() {
+			idx := 0
+
+			for item := range seq {
+				jobs <- work{idx, item}
+
+				idx++
+			}
+
+			close(jobs)
+
+			wg.Wait()
+
+			close(results)
+		}()
+
+		next := 0
+
+		pending := make(map[int]Out)
+
+		for r := range results {
+			pending[r.idx] = r.out
+
+			for {
+				v, ok := pending[next]
+
+				if !ok {
+					break
+				}
+
+				out <- v
+
+				delete(pending, next)
+
+				next++
+			}
+		}
+	}()
+
+	return out
+}