Support dynamic select (#7)

* Support dynamic select * Can take various sequences * Rename to `any` from `forEach` * Add an overload that can take variadic arguments
gh123man · Jul 7, 2024 · 394b971 · 394b971
1 parent 679ee7d
commit 394b971
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -134,6 +134,8 @@ The loop will break when the channel is closed.
 
 `none { ... }` if none of the channel operations were ready, none will execute instead. 
 
+`any(x1, x2, ...) { x in ... }` or `any(seq) { el in ... }` operates on a sequence and is useful for working with an array of channels.
+
 ### Examples
 
 ```swift

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,48 @@ 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 any<S, T>(_ seq: S, @SelectCollector cases: (T) -> ([SelectHandler])) -> [SelectHandler] where S: Sequence, S.Element == T {
+    return seq.flatMap { cases($0) }
+}
+
+@inlinable
+@inline(__always)
+public func any<T>(_ elements: T..., @SelectCollector cases: (T) -> ([SelectHandler])) -> [SelectHandler] {
+    return elements.flatMap { cases($0) }
+}
diff --git a/Tests/AsyncChannelsTests/AsyncChannelTests.swift b/Tests/AsyncChannelsTests/AsyncChannelTests.swift
@@ -212,7 +212,61 @@ 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 {
+            any([a, b]) {
+                receive($0) { await result <- $0! }
+            }
+        }
+
+        await select {
+            any([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 testDynamicVariadicSelect() async {
+    let a = Channel<String>()
+    let b = Channel<String>()
+    let result = Channel<String>(capacity: 2)
+
+    Task {
+        await a <- "foo"
+        await b <- "bar"
+    }
+
+    await select {
+        any(a, b) {
+            receive($0) { await result <- $0! }
+        }
+    }
+
+    await select {
+        any(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)