Add WithSplitFunc option (WithTimeout now need generic, damn)

Author: floatdrop

batchan.go

@@ -2,29 +2,40 @@ package batchan
 
 import "time"
 
-type option func(*config)
+type option[T any] func(*config[T])
 
-type config struct {
+type config[T any] struct {
 	timeout    time.Duration
 	hasTimeout bool
+	splitFunc  func(T, T) bool
 }
 
-func WithTimeout(timeout time.Duration) option {
-	return func(cfg *config) {
+func WithTimeout[T any](timeout time.Duration) option[T] {
+	return func(cfg *config[T]) {
 		cfg.timeout = timeout
 		cfg.hasTimeout = true
 	}
 }
 
+func WithSplitFunc[T any](splitFunc func(T, T) bool) option[T] {
+	return func(cfg *config[T]) {
+		cfg.splitFunc = splitFunc
+	}
+}
+
 func timerOrNil(t *time.Timer) <-chan time.Time {
 	if t != nil {
 		return t.C
 	}
 	return nil
 }
 
-func New[T any](in <-chan T, size int, opts ...option) <-chan []T {
-	cfg := &config{}
+func noSplitFunc[T any](t1, t2 T) bool { return false }
+
+func New[T any](in <-chan T, size int, opts ...option[T]) <-chan []T {
+	cfg := &config[T]{
+		splitFunc: noSplitFunc[T],
+	}
 
 	for _, opt := range opts {
 		opt(cfg)
@@ -63,11 +74,14 @@ func New[T any](in <-chan T, size int, opts ...option) <-chan []T {
 					return
 				}
 
+				if len(currentBatch) > 0 && cfg.splitFunc(currentBatch[len(currentBatch)-1], t) {
+					flush()
+				}
+
 				currentBatch = append(currentBatch, t)
 				if len(currentBatch) >= size {
 					flush()
 				}
-
 			case <-timerOrNil(timer):
 				flush()
 			}

batchan_test.go

@@ -30,7 +30,7 @@ func collectBatches[T any](out <-chan []T) [][]T {
 
 func TestBatchingFlushOnTimeout(t *testing.T) {
 	in := sendIntsToChan([]int{1, 2}, 150*time.Millisecond) // delay > timeout
-	out := batchan.New(in, 5, batchan.WithTimeout(100*time.Millisecond))
+	out := batchan.New(in, 5, batchan.WithTimeout[int](100*time.Millisecond))
 
 	got := collectBatches(out)
 
@@ -43,7 +43,7 @@ func TestBatchingFlushOnTimeout(t *testing.T) {
 
 func TestBatchingFlushOnSizeOrTimeout(t *testing.T) {
 	in := sendIntsToChan([]int{1, 2, 3, 4}, 50*time.Millisecond)
-	out := batchan.New(in, 2, batchan.WithTimeout(200*time.Millisecond))
+	out := batchan.New(in, 2, batchan.WithTimeout[int](200*time.Millisecond))
 
 	got := collectBatches(out)
 	expected := [][]int{{1, 2}, {3, 4}}
@@ -55,7 +55,7 @@ func TestBatchingFlushOnSizeOrTimeout(t *testing.T) {
 
 func TestFlushTimeoutMultiple(t *testing.T) {
 	in := sendIntsToChan([]int{1, 2, 3}, 300*time.Millisecond)
-	out := batchan.New(in, 10, batchan.WithTimeout(200*time.Millisecond)) // small timeout, large batch size
+	out := batchan.New(in, 10, batchan.WithTimeout[int](200*time.Millisecond)) // small timeout, large batch size
 
 	got := collectBatches(out)
 	expected := [][]int{{1}, {2}, {3}}
@@ -77,7 +77,7 @@ func TestTimeoutResetsAfterFlush(t *testing.T) {
 		in <- 3
 	}()
 
-	out := batchan.New(in, 2, batchan.WithTimeout(100*time.Millisecond))
+	out := batchan.New(in, 2, batchan.WithTimeout[int](100*time.Millisecond))
 
 	got := collectBatches(out)
 	expected := [][]int{{1}, {2}, {3}}
@@ -146,6 +146,18 @@ func TestBatchSizeLargerThanInput(t *testing.T) {
 	}
 }
 
+func TestSplitFunc(t *testing.T) {
+	in := sendIntsToChan([]int{1, 2, 3, 5, 6}, time.Microsecond)
+	out := batchan.New(in, 5, batchan.WithSplitFunc(func(i1, i2 int) bool { return i2-i1 > 1 }))
+
+	expected := [][]int{{1, 2, 3}, {5, 6}}
+	got := collectBatches(out)
+
+	if !reflect.DeepEqual(got, expected) {
+		t.Errorf("Expected %v, got %v", expected, got)
+	}
+}
+
 // Optional: Test that the output channel closes properly
 func TestOutputChannelClosure(t *testing.T) {
 	in := sendIntsToChan([]int{1, 2, 3}, time.Microsecond)