Support dynamic select

gh123man · Jul 7, 2024 · f5f51ad · f5f51ad
1 parent 679ee7d
commit f5f51ad
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 15 deletions.
diff --git a/Sources/AsyncChannels/Select.swift b/Sources/AsyncChannels/Select.swift
@@ -83,8 +83,8 @@ struct SendHandler<T>: SelectProtocol {
 
 @resultBuilder
 public struct SelectCollector {
-    public static func buildBlock(_ handlers: SelectHandler...) -> [SelectHandler] {
-        return handlers
+    public static func buildBlock(_ handlers: [SelectHandler]...) -> [SelectHandler] {
+        return handlers.reduce([], +)
     }
 }
 
@@ -118,37 +118,43 @@ public func select(@SelectCollector cases: () -> ([SelectHandler])) async {
 
 @inlinable
 @inline(__always)
-public func receive<T>(_ chan: Channel<T>, _ outFunc: @escaping (T?) async -> ()) -> SelectHandler {
-    return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: outFunc))
+public func receive<T>(_ chan: Channel<T>, _ outFunc: @escaping (T?) async -> ()) -> [SelectHandler] {
+    return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: outFunc))]
 }
 
 @inlinable
 @inline(__always)
-public func receive<T>(_ chan: Channel<T>, _ outFunc: @escaping () async -> ()) -> SelectHandler {
-    return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in await outFunc() }))
+public func receive<T>(_ chan: Channel<T>, _ outFunc: @escaping () async -> ()) -> [SelectHandler] {
+    return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in await outFunc() }))]
 }
 
 @inlinable
 @inline(__always)
-public func receive<T>(_ chan: Channel<T>) -> SelectHandler {
-    return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in }))
+public func receive<T>(_ chan: Channel<T>) -> [SelectHandler] {
+    return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in }))]
 }
 
 @inlinable
 @inline(__always)
-public func send<T>(_ val: T, to chan: Channel<T>) -> SelectHandler {
-    return SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: {}))
+public func send<T>(_ val: T, to chan: Channel<T>) -> [SelectHandler] {
+    return [SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: {}))]
 }
 
 @inlinable
 @inline(__always)
-public func send<T>(_ val: T, to chan: Channel<T>, _ onSend: @escaping () async -> ()) -> SelectHandler {
-    return SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: onSend))
+public func send<T>(_ val: T, to chan: Channel<T>, _ onSend: @escaping () async -> ()) -> [SelectHandler] {
+    return [SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: onSend))]
 }
 
 @inlinable
 @inline(__always)
-public func none(handler: @escaping () async -> ()) -> SelectHandler {
-    return SelectHandler(inner: NoneHandler(handler: handler))
+public func none(handler: @escaping () async -> ()) -> [SelectHandler] {
+    return [SelectHandler(inner: NoneHandler(handler: handler))]
+}
+
+@inlinable
+@inline(__always)
+public func forEach<T>(_ chans: [Channel<T>], @SelectCollector cases: (Channel<T>) -> ([SelectHandler])) -> [SelectHandler] {
+    return chans.flatMap { cases($0) }
 }
 
diff --git a/Tests/AsyncChannelsTests/AsyncChannelTests.swift b/Tests/AsyncChannelsTests/AsyncChannelTests.swift
@@ -212,7 +212,34 @@ final class AsyncTest: XCTestCase {
         let r = await result.reduce(into: []) { $0.append($1) }
         XCTAssertEqual(["foo", "bar"].sorted(), r.sorted())
     }
-
+
+    func testDynamicSelect() async {
+        let a = Channel<String>()
+        let b = Channel<String>()
+        let result = Channel<String>(capacity: 2)
+
+        Task {
+            await a <- "foo"
+            await b <- "bar"
+        }
+
+        await select {
+            forEach([a, b]) {
+                receive($0) { await result <- $0! }
+            }
+        }
+
+        await select {
+            forEach([a, b]) {
+                receive($0) { await result <- $0! }
+            }
+        }
+        result.close()
+
+        let r = await result.reduce(into: []) { $0.append($1) }
+        XCTAssertEqual(["foo", "bar"].sorted(), r.sorted())
+    }
+
     func testBufferSelect() async {
         let c = Channel<String>(capacity: 3)
         let d = Channel<String>(capacity: 3)