From 369755e88f658123e9044acb109044728444d20e Mon Sep 17 00:00:00 2001
From: Daniel Alm <Daniel.Alm@ForumD.net>
Date: Mon, 1 Apr 2019 11:05:18 +0200
Subject: [PATCH] Significantly improve the performance of
 `Collection<FutureType>.flatten`.

This further halves the runtime in https://github.com/vapor/template-kit/pull/50 from 0.1s to 50ms.
---
 Sources/Async/Future+Flatten.swift | 47 +++++++++++++++++++-----------
 Sources/Async/FutureType.swift     |  3 ++
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/Sources/Async/Future+Flatten.swift b/Sources/Async/Future+Flatten.swift
index c4db4d01..be53919e 100644
--- a/Sources/Async/Future+Flatten.swift
+++ b/Sources/Async/Future+Flatten.swift
@@ -60,27 +60,40 @@ extension Collection where Element: FutureType {
             return eventLoop.newSucceededFuture(result: [])
         }
         
-        var promises = [EventLoopPromise<Element.Expectation>]()
-        for future in self {
-            let promise = eventLoop.newPromise(of: Element.Expectation.self)
-            promises.append(promise)
+        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)
+        for (index, future) in self.enumerated() {
             future.addAwaiter { result in
-                switch result {
-                case .success(let value):
-                    promise.succeed(result: value)
-                case .error(let error):
-                    promise.fail(error: error)
+                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)
                 }
             }
         }
-        let futures = promises.map { $0.futureResult }
-        return Future<[Element.Expectation]>.reduce(
-            into: [],
-            futures,
-            eventLoop: eventLoop
-        ) { partialResult, nextElement in
-            return partialResult.append(nextElement)
-        }
+        return resultPromise.futureResult
     }
 }
 
diff --git a/Sources/Async/FutureType.swift b/Sources/Async/FutureType.swift
index 8a698832..af62ebcf 100644
--- a/Sources/Async/FutureType.swift
+++ b/Sources/Async/FutureType.swift
@@ -10,6 +10,9 @@ public protocol FutureType {
     /// This future's result type.
     typealias Result = FutureResult<Expectation>
 
+    /// The event loop this future is fulfilled on.
+    var eventLoop: EventLoop { get }
+
     /// Adds a new awaiter to this `Future` that will be called when the result is ready.
     func addAwaiter(callback: @escaping FutureResultCallback<Expectation>)
 }