diff --git a/Sources/AsyncChannels/Select.swift b/Sources/AsyncChannels/Select.swift index 9f98944..9b03f88 100644 --- a/Sources/AsyncChannels/Select.swift +++ b/Sources/AsyncChannels/Select.swift @@ -83,8 +83,8 @@ struct SendHandler: 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(_ chan: Channel, _ outFunc: @escaping (T?) async -> ()) -> SelectHandler { - return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: outFunc)) +public func receive(_ chan: Channel, _ outFunc: @escaping (T?) async -> ()) -> [SelectHandler] { + return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: outFunc))] } @inlinable @inline(__always) -public func receive(_ chan: Channel, _ outFunc: @escaping () async -> ()) -> SelectHandler { - return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in await outFunc() })) +public func receive(_ chan: Channel, _ outFunc: @escaping () async -> ()) -> [SelectHandler] { + return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in await outFunc() }))] } @inlinable @inline(__always) -public func receive(_ chan: Channel) -> SelectHandler { - return SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in })) +public func receive(_ chan: Channel) -> [SelectHandler] { + return [SelectHandler(inner: ReceiveHandler(chan: chan, outFunc: { _ in }))] } @inlinable @inline(__always) -public func send(_ val: T, to chan: Channel) -> SelectHandler { - return SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: {})) +public func send(_ val: T, to chan: Channel) -> [SelectHandler] { + return [SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: {}))] } @inlinable @inline(__always) -public func send(_ val: T, to chan: Channel, _ onSend: @escaping () async -> ()) -> SelectHandler { - return SelectHandler(inner: SendHandler(chan: chan, val: val, onSend: onSend)) +public func send(_ val: T, to chan: Channel, _ 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(_ chans: [Channel], @SelectCollector cases: (Channel) -> ([SelectHandler])) -> [SelectHandler] { + return chans.flatMap { cases($0) } } diff --git a/Tests/AsyncChannelsTests/AsyncChannelTests.swift b/Tests/AsyncChannelsTests/AsyncChannelTests.swift index 96ac4ed..6416ec4 100644 --- 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() + let b = Channel() + let result = Channel(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(capacity: 3) let d = Channel(capacity: 3)