feat: add Val{N} parameter to tee for N-way channel splitting

src/channel_utilities.jl

@@ -29,47 +29,64 @@ end
 
 """
     tee(in::AbstractChannel, channel_size::Integer=16)
+    tee(::Val{N}, in::AbstractChannel, channel_size::Integer=16) where N
 
-Split a channel into two synchronized outputs. Both output channels receive
+Split a channel into N synchronized outputs. All output channels receive
 identical copies of the data.
 
-Returns a tuple `(out1, out2)` of two channels with the same type as the input.
+Returns a tuple of N channels with the same type as the input.
 
 For `SignalChannel`, preserves the matrix dimensions.
 For generic `PipeChannel`, creates output channels with the specified buffer size.
 
 Based on benchmarks in benchmark/benchmarks.jl a channel size of 16 is a sweet spot.
 
 # Arguments
+- `Val{N}`: Number of output channels (default: 2 when not specified)
 - `in`: Input channel to split
 - `channel_size`: Buffer size for output channels (default: 16)
 
 # Examples
 ```julia
-# With SignalChannel
+# With SignalChannel (default 2 outputs)
 input = SignalChannel{ComplexF32}(1024, 4)
 out1, out2 = tee(input, 16)
 
 # With generic PipeChannel
 input = PipeChannel{Int}(10)
 out1, out2 = tee(input, 16)
+
+# With 3 outputs
+input = SignalChannel{ComplexF32}(1024, 4)
+out1, out2, out3 = tee(Val(3), input, 16)
+
+# With 4 outputs
+input = PipeChannel{Int}(10)
+outputs = tee(Val(4), input)  # Returns NTuple{4, PipeChannel{Int}}
 ```
 """
-function tee(in::AbstractChannel, channel_size::Integer=16)
-    out1 = similar(in, channel_size)
-    out2 = similar(in, channel_size)
+function tee(::Val{N}, in::AbstractChannel, channel_size::Integer=16) where N
+    outputs = ntuple(_ -> similar(in, channel_size), Val(N))
     task = Threads.@spawn begin
         for data in in
-            put!(out1, data)
-            put!(out2, data)
+            for out in outputs
+                put!(out, data)
+            end
         end
-        close(out1)
-        close(out2)
+        for out in outputs
+            close(out)
+        end
+    end
+    for out in outputs
+        bind(out, task)
     end
-    bind(out1, task)
-    bind(out2, task)
     bind(in, task)
-    return (out1, out2)
+    return outputs
+end
+
+# Default to 2 outputs for backward compatibility
+function tee(in::AbstractChannel, channel_size::Integer=16)
+    return tee(Val(2), in, channel_size)
 end
 
 """

test/channel_utilities.jl

@@ -113,6 +113,85 @@ using FixedSizeArrays: FixedSizeMatrixDefault
         @test results2 == [1, 2, 3, 4, 5]
     end
 
+    @testset "tee with Val{3}" begin
+        input_chan = PipeChannel{Int}(10)
+        out1, out2, out3 = tee(Val(3), input_chan)
+
+        task = @async begin
+            for i in 1:5
+                put!(input_chan, i)
+            end
+            close(input_chan)
+        end
+
+        results1 = []
+        results2 = []
+        results3 = []
+
+        # Consume all channels concurrently
+        task1 = @async begin
+            for data in out1
+                push!(results1, data)
+            end
+        end
+
+        task2 = @async begin
+            for data in out2
+                push!(results2, data)
+            end
+        end
+
+        task3 = @async begin
+            for data in out3
+                push!(results3, data)
+            end
+        end
+
+        wait(task)
+        wait(task1)
+        wait(task2)
+        wait(task3)
+        @test results1 == [1, 2, 3, 4, 5]
+        @test results2 == [1, 2, 3, 4, 5]
+        @test results3 == [1, 2, 3, 4, 5]
+    end
+
+    @testset "tee with Val{4} and SignalChannel" begin
+        input_chan = SignalChannel{ComplexF32,2}(100)
+        outputs = tee(Val(4), input_chan)
+
+        @test length(outputs) == 4
+
+        task = @async begin
+            for i in 1:3
+                data = FixedSizeMatrixDefault{ComplexF32}(fill(ComplexF32(i, 0), 100, 2))
+                put!(input_chan, data)
+            end
+            close(input_chan)
+        end
+
+        all_results = [[] for _ in 1:4]
+
+        # Consume all channels concurrently
+        consumer_tasks = map(enumerate(outputs)) do (idx, out)
+            @async begin
+                for data in out
+                    push!(all_results[idx], data[1, 1])
+                end
+            end
+        end
+
+        wait(task)
+        for t in consumer_tasks
+            wait(t)
+        end
+
+        expected = [ComplexF32(1, 0), ComplexF32(2, 0), ComplexF32(3, 0)]
+        for results in all_results
+            @test results == expected
+        end
+    end
+
     @testset "write_to_file" begin
         mktempdir() do tmpdir
             input_chan = SignalChannel{ComplexF32,3}(100, 10)  # Buffer size to prevent blocking