-
-
Save wilg/47a04c8f5083a6938da6087f77333784 to your computer and use it in GitHub Desktop.
Swift async/await implementation of a parallel map
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Foundation | |
// https://gist.github.com/DougGregor/92a2e4f6e11f6d733fb5065e9d1c880f | |
extension Collection { | |
func parallelMap<T>( | |
parallelism requestedParallelism: Int? = nil, | |
_ transform: @escaping (Element) async throws -> T | |
) async rethrows -> [T] { | |
let defaultParallelism = 2 | |
let parallelism = requestedParallelism ?? defaultParallelism | |
let n = count | |
if n == 0 { | |
return [] | |
} | |
return try await withThrowingTaskGroup(of: (Int, T).self, returning: [T].self) { group in | |
var result = [T?](repeatElement(nil, count: n)) | |
var i = self.startIndex | |
var submitted = 0 | |
func submitNext() async throws { | |
if i == self.endIndex { return } | |
group.addTask { [submitted, i] in | |
let value = try await transform(self[i]) | |
return (submitted, value) | |
} | |
submitted += 1 | |
formIndex(after: &i) | |
} | |
// submit first initial tasks | |
for _ in 0 ..< parallelism { | |
try await submitNext() | |
} | |
// as each task completes, submit a new task until we run out of work | |
while let (index, taskResult) = try await group.next() { | |
result[index] = taskResult | |
try Task.checkCancellation() | |
try await submitNext() | |
} | |
assert(result.count == n) | |
return Array(result.compactMap { $0 }) | |
} | |
} | |
func parallelEach( | |
parallelism requestedParallelism: Int? = nil, | |
_ work: @escaping (Element) async throws -> Void | |
) async rethrows { | |
_ = try await parallelMap { | |
try await work($0) | |
} | |
} | |
} |
Line 52 should be
_ = try await parallelMap(parallelism: parallelism) {
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Is this still the recommended way to do an async map or does the Swift Async Algorithms package have something better?
(also just curious why is this was not a part of Foundation originally -- seems like such an obvious need)