Skip to content

Instantly share code, notes, and snippets.

@adam-fowler
Last active November 17, 2021 16:08
Show Gist options
  • Save adam-fowler/d94040ed636db975812c4e5742cbeaa6 to your computer and use it in GitHub Desktop.
Save adam-fowler/d94040ed636db975812c4e5742cbeaa6 to your computer and use it in GitHub Desktop.
Sequence `forEach` and `map` that uses an async closure

Implement Sequence.forEach and Sequence.map that takes an async closure.

@available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *)
extension Sequence {
    // named `asyncForEach` as naming it `forEach` confuses the compiler
    func asyncForEach(_ body: @escaping (Element) async throws -> Void) async rethrows {
        try await withThrowingTaskGroup(of: Void.self) { group in
            self.forEach { element in
                group.addTask {
                    try await body(element)
                }
            }
            try await group.waitForAll()
        }
    }
}

@available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *)
extension Sequence {
    public func asyncMap<T: Sendable>(_ transform: @escaping (Element) async throws -> T) async rethrows -> [T] {
        let result: ContiguousArray<(Int, T)> = try await withThrowingTaskGroup(of: (Int, T).self) { group in
            self.enumerated().forEach { element in
                group.addTask {
                    let result = try await transform(element.1)
                    return (element.0, result)
                }
            }
            // Code for collating results copied from Sequence.map in Swift codebase
            let initialCapacity = underestimatedCount
            var result = ContiguousArray<(Int, T)>()
            result.reserveCapacity(initialCapacity)

            // Add elements up to the initial capacity without checking for regrowth.
            for _ in 0..<initialCapacity {
                try await result.append(group.next()!)
            }
            // Add remaining elements, if any.
            while let element = try await group.next() {
                result.append(element)
            }
            return result
        }
        // construct final array and fill in elements
        return Array<T>(unsafeUninitializedCapacity: result.count) { buffer, count in
            for value in result {
                (buffer.baseAddress! + value.0).initialize(to: value.1)
            }
            count = result.count
        }
    }
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment