feat(iterator): unoptimized Tee

Author: shils

iterator/iter.go

@@ -2,6 +2,8 @@ package iterator
 
 import (
 	"iter"
+	"slices"
+	"sync"
 )
 
 func Map[T, V any](it func(func(T) bool), fn func(T) V) func(func(V) bool) {
@@ -220,3 +222,64 @@ func Flatten[T any, S ~[]T](it func(func(S) bool)) func(func(T) bool) {
 		}
 	}
 }
+
+type teeState[T any] struct {
+	next      func() (T, bool)
+	mu        sync.Mutex
+	buf       []T
+	positions []int
+}
+
+func (s *teeState[T]) advanceOne(i int) (T, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.positions[i] == len(s.buf) {
+		t, ok := s.next()
+		if !ok {
+			return t, ok
+		}
+		s.buf = append(s.buf, t)
+		minPos := slices.Min(s.positions)
+		if minPos > 0 {
+			s.buf = s.buf[minPos:]
+			for j := range s.positions {
+				s.positions[j] -= minPos
+			}
+		}
+	}
+	pos := s.positions[i]
+	s.positions[i]++
+	return s.buf[pos], true
+}
+
+func Tee[T any](it func(func(T) bool), n int) []func(func(T) bool) {
+	next, stop := iter.Pull(it)
+
+	state := &teeState[T]{
+		next:      next,
+		positions: make([]int, n),
+	}
+	stopped := 0
+
+	outs := make([]func(func(T) bool), n)
+	for i := range n {
+		i := i
+		outs[i] = func(yield func(T) bool) {
+			for {
+				t, ok := state.advanceOne(i)
+				if !ok {
+					return
+				}
+				if !yield(t) {
+					stopped++
+					break
+				}
+			}
+			if stopped == n {
+				stop()
+			}
+		}
+	}
+	return outs
+}

iterator/iter_test.go

@@ -1,11 +1,13 @@
 package iterator
 
 import (
+	"go-exp/functions/partials"
 	"gotest.tools/v3/assert"
 	"maps"
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -198,3 +200,43 @@ func TestFlatten(t *testing.T) {
 
 	assert.DeepEqual(t, result, expected)
 }
+
+func TestTee(t *testing.T) {
+	it := slices.Values([]int{4, 3, 2, 1})
+	its := Tee(it, 2)
+
+	result1 := slices.Collect(its[0])
+	result2 := slices.Collect(its[1])
+
+	expected := []int{4, 3, 2, 1}
+	assert.DeepEqual(t, result1, expected)
+	assert.DeepEqual(t, result2, expected)
+}
+
+func BenchmarkTee(b *testing.B) {
+	if b.N > 10000000 {
+		b.Skipf("N too large: %d", b.N)
+	}
+
+	it := TakeWhile(Generate(b.N, func(x int) int { return x - 1 }), partials.Gt(0))
+	its := Tee(it, 10)
+
+	results := make([][]int, 10)
+
+	var wg sync.WaitGroup
+
+	for i := range 10 {
+		wg.Add(1)
+		go func() {
+			results[i] = slices.Collect(its[i])
+			wg.Done()
+		}()
+	}
+
+	expected := slices.Collect(TakeWhile(Generate(b.N, func(x int) int { return x - 1 }), partials.Gt(0)))
+
+	wg.Wait()
+	for _, result := range results {
+		assert.DeepEqual(b, result, expected)
+	}
+}