1 files changed, 60 insertions, 10 deletions

diff --git a/llvm/include/llvm/Support/ThreadPool.h b/llvm/include/llvm/Support/ThreadPool.h
index 4c41b88d6043..8d30e8e92755 100644
--- a/llvm/include/llvm/Support/ThreadPool.h
+++ b/llvm/include/llvm/Support/ThreadPool.h
@@ -36,9 +36,6 @@ namespace llvm {
 /// for some work to become available.
 class ThreadPool {
 public:
-  using TaskTy = std::function<void()>;
-  using PackagedTaskTy = std::packaged_task<void()>;
-
   /// Construct a pool using the hardware strategy \p S for mapping hardware
   /// execution resources (threads, cores, CPUs)
   /// Defaults to using the maximum execution resources in the system, but
@@ -51,17 +48,17 @@ public:
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
   template <typename Function, typename... Args>
-  inline std::shared_future<void> async(Function &&F, Args &&... ArgList) {
+  inline auto async(Function &&F, Args &&...ArgList) {
     auto Task =
         std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
-    return asyncImpl(std::move(Task));
+    return async(std::move(Task));
   }
 
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
-  template <typename Function>
-  inline std::shared_future<void> async(Function &&F) {
-    return asyncImpl(std::forward<Function>(F));
+  template <typename Func>
+  auto async(Func &&F) -> std::shared_future<decltype(F())> {
+    return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)));
   }
 
   /// Blocking wait for all the threads to complete and the queue to be empty.
@@ -74,17 +71,70 @@ public:
   bool isWorkerThread() const;
 
 private:
+  /// Helpers to create a promise and a callable wrapper of \p Task that sets
+  /// the result of the promise. Returns the callable and a future to access the
+  /// result.
+  template <typename ResTy>
+  static std::pair<std::function<void()>, std::future<ResTy>>
+  createTaskAndFuture(std::function<ResTy()> Task) {
+    std::shared_ptr<std::promise<ResTy>> Promise =
+        std::make_shared<std::promise<ResTy>>();
+    auto F = Promise->get_future();
+    return {
+        [Promise = std::move(Promise), Task]() { Promise->set_value(Task()); },
+        std::move(F)};
+  }
+  static std::pair<std::function<void()>, std::future<void>>
+  createTaskAndFuture(std::function<void()> Task) {
+    std::shared_ptr<std::promise<void>> Promise =
+        std::make_shared<std::promise<void>>();
+    auto F = Promise->get_future();
+    return {[Promise = std::move(Promise), Task]() {
+              Task();
+              Promise->set_value();
+            },
+            std::move(F)};
+  }
+
   bool workCompletedUnlocked() { return !ActiveThreads && Tasks.empty(); }
 
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
-  std::shared_future<void> asyncImpl(TaskTy F);
+  template <typename ResTy>
+  std::shared_future<ResTy> asyncImpl(std::function<ResTy()> Task) {
+
+#if LLVM_ENABLE_THREADS
+    /// Wrap the Task in a std::function<void()> that sets the result of the
+    /// corresponding future.
+    auto R = createTaskAndFuture(Task);
+
+    {
+      // Lock the queue and push the new task
+      std::unique_lock<std::mutex> LockGuard(QueueLock);
+
+      // Don't allow enqueueing after disabling the pool
+      assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
+      Tasks.push(std::move(R.first));
+    }
+    QueueCondition.notify_one();
+    return R.second.share();
+
+#else // LLVM_ENABLE_THREADS Disabled
+
+    // Get a Future with launch::deferred execution using std::async
+    auto Future = std::async(std::launch::deferred, std::move(Task)).share();
+    // Wrap the future so that both ThreadPool::wait() can operate and the
+    // returned future can be sync'ed on.
+    Tasks.push([Future]() { Future.get(); });
+    return Future;
+#endif
+  }
 
   /// Threads in flight
   std::vector<llvm::thread> Threads;
 
   /// Tasks waiting for execution in the pool.
-  std::queue<PackagedTaskTy> Tasks;
+  std::queue<std::function<void()>> Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;