Flatten.swift

extension Collection where Element: FutureType {
    /// Flattens an array of futures into a future with an array of results.
    /// - note: the order of the results will match the order of the futures in the input array.
    public func flatten(on worker: Worker) -> Future<[Element.Expectation]> {
        let eventLoop = worker.eventLoop
        
        // Avoid unnecessary work
        guard count > 0 else {
            return eventLoop.newSucceededFuture(result: [])
        }
        
        let resultPromise: EventLoopPromise<[Element.Expectation]> = eventLoop.newPromise()
        var promiseFulfilled = false
        
        let expectedCount = self.count
        var fulfilledCount = 0
        var results = Array<Element.Expectation?>(repeating: nil, count: expectedCount)
        let inEventLoop = eventLoop.inEventLoop
        for (index, future) in self.enumerated() {
            if inEventLoop && future.eventLoop === eventLoop,
                let result = (future as? Future<Element.Expectation>)?.value {
                switch result {
                case .success(let result):
                    results[index] = result
                    fulfilledCount += 1

                    if fulfilledCount == expectedCount {
                        promiseFulfilled = true
                        // Forcibly unwrapping is okay here, because we know that each result has been filled.
                        resultPromise.succeed(result: results.map { $0! })
                    }
                case .failure(let error):
                    promiseFulfilled = true
                    resultPromise.fail(error: error)
                }
            } else {
                future.addAwaiter { result in
                    let work: () -> Void = {
                        guard !promiseFulfilled else { return }
                        switch result {
                        case .success(let result):
                            results[index] = result
                            fulfilledCount += 1

                            if fulfilledCount == expectedCount {
                                promiseFulfilled = true
                                // Forcibly unwrapping is okay here, because we know that each result has been filled.
                                resultPromise.succeed(result: results.map { $0! })
                            }
                        case .error(let error):
                            promiseFulfilled = true
                            resultPromise.fail(error: error)
                        }
                    }

                    if future.eventLoop === eventLoop {
                        work()
                    } else {
                        eventLoop.execute(work)
                    }
                }
            }
        }
        return resultPromise.futureResult
    }
}