From 753868d5ecf3a89595e4a49a1389e5b2b9dd4b9f Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Mon, 15 Feb 2016 11:32:24 -0800
Subject: [PATCH] Unify runInMainContext for void and non-void

Summary: This also fixes a bug where exception was not re-thrown for functions returning void.

Reviewed By: spalamarchuk

Differential Revision: D2936887

fb-gh-sync-id: 9828dec131203528c27eae874aba147168f40d0d
shipit-source-id: 9828dec131203528c27eae874aba147168f40d0d
---
 folly/experimental/fibers/FiberManager-inl.h  | 24 +----------------
 folly/experimental/fibers/FiberManager.h      | 18 -------------
 folly/experimental/fibers/test/FibersTest.cpp | 26 ++++++++++++-------
 folly/futures/Try-inl.h                       | 10 +++++--
 folly/futures/Try.h                           | 11 ++++++--
 5 files changed, 34 insertions(+), 55 deletions(-)

diff --git a/folly/experimental/fibers/FiberManager-inl.h b/folly/experimental/fibers/FiberManager-inl.h
index 8c22010b..b0a6334e 100644
--- a/folly/experimental/fibers/FiberManager-inl.h
+++ b/folly/experimental/fibers/FiberManager-inl.h
@@ -429,14 +429,6 @@ void FiberManager::addTaskFinally(F&& func, G&& finally) {
 template <typename F>
 typename std::result_of<F()>::type
 FiberManager::runInMainContext(F&& func) {
-  return runInMainContextHelper(std::forward<F>(func));
-}
-
-template <typename F>
-inline typename std::enable_if<
-  !std::is_same<typename std::result_of<F()>::type, void>::value,
-  typename std::result_of<F()>::type>::type
-FiberManager::runInMainContextHelper(F&& func) {
   if (UNLIKELY(activeFiber_ == nullptr)) {
     return func();
   }
@@ -451,21 +443,7 @@ FiberManager::runInMainContextHelper(F&& func) {
   immediateFunc_ = std::ref(f);
   activeFiber_->preempt(Fiber::AWAITING_IMMEDIATE);
 
-  return std::move(result.value());
-}
-
-template <typename F>
-inline typename std::enable_if<
-  std::is_same<typename std::result_of<F()>::type, void>::value,
-  void>::type
-FiberManager::runInMainContextHelper(F&& func) {
-  if (UNLIKELY(activeFiber_ == nullptr)) {
-    func();
-    return;
-  }
-
-  immediateFunc_ = std::ref(func);
-  activeFiber_->preempt(Fiber::AWAITING_IMMEDIATE);
+  return std::move(result).value();
 }
 
 inline FiberManager& FiberManager::getFiberManager() {
diff --git a/folly/experimental/fibers/FiberManager.h b/folly/experimental/fibers/FiberManager.h
index 92d9cef3..f90b9d67 100644
--- a/folly/experimental/fibers/FiberManager.h
+++ b/folly/experimental/fibers/FiberManager.h
@@ -373,24 +373,6 @@ class FiberManager : public ::folly::Executor {
    */
   static FOLLY_TLS FiberManager* currentFiberManager_;
 
-  /**
-   * runInMainContext implementation for non-void functions.
-   */
-  template <typename F>
-  typename std::enable_if<
-    !std::is_same<typename std::result_of<F()>::type, void>::value,
-    typename std::result_of<F()>::type>::type
-  runInMainContextHelper(F&& func);
-
-  /**
-   * runInMainContext implementation for void functions
-   */
-  template <typename F>
-  typename std::enable_if<
-    std::is_same<typename std::result_of<F()>::type, void>::value,
-    void>::type
-  runInMainContextHelper(F&& func);
-
   /**
    * Allocator used to allocate stack for Fibers in the pool.
    * Allocates stack on the stack of the main context.
diff --git a/folly/experimental/fibers/test/FibersTest.cpp b/folly/experimental/fibers/test/FibersTest.cpp
index e3eb2f74..b7c69c9c 100644
--- a/folly/experimental/fibers/test/FibersTest.cpp
+++ b/folly/experimental/fibers/test/FibersTest.cpp
@@ -937,16 +937,22 @@ TEST(FiberManager, runInMainContext) {
 
   checkRan = false;
 
-  manager.addTask(
-    [&]() {
-      int stackLocation;
-      runInMainContext(
-        [&]() {
-          expectMainContext(checkRan, &mainLocation, &stackLocation);
-        });
-      EXPECT_TRUE(checkRan);
-    }
-  );
+  manager.addTask([&]() {
+    struct A {
+      explicit A(int value_) : value(value_) {}
+      A(const A&) = delete;
+      A(A&&) = default;
+
+      int value;
+    };
+    int stackLocation;
+    auto ret = runInMainContext([&]() {
+      expectMainContext(checkRan, &mainLocation, &stackLocation);
+      return A(42);
+    });
+    EXPECT_TRUE(checkRan);
+    EXPECT_EQ(42, ret.value);
+  });
 
   loopController.loop(
     [&]() {
diff --git a/folly/futures/Try-inl.h b/folly/futures/Try-inl.h
index 1d0d1d27..499dfa44 100644
--- a/folly/futures/Try-inl.h
+++ b/folly/futures/Try-inl.h
@@ -102,13 +102,19 @@ Try<T>::~Try() {
 }
 
 template <class T>
-T& Try<T>::value() {
+T& Try<T>::value() & {
   throwIfFailed();
   return value_;
 }
 
 template <class T>
-const T& Try<T>::value() const {
+T&& Try<T>::value() && {
+  throwIfFailed();
+  return std::move(value_);
+}
+
+template <class T>
+const T& Try<T>::value() const & {
   throwIfFailed();
   return value_;
 }
diff --git a/folly/futures/Try.h b/folly/futures/Try.h
index 437127b8..b4152dd4 100644
--- a/folly/futures/Try.h
+++ b/folly/futures/Try.h
@@ -124,14 +124,21 @@ class Try {
    *
    * @returns mutable reference to the contained value
    */
-  T& value();
+  T& value()&;
+  /*
+   * Get a rvalue reference to the contained value. If the Try contains an
+   * exception it will be rethrown.
+   *
+   * @returns rvalue reference to the contained value
+   */
+  T&& value()&&;
   /*
    * Get a const reference to the contained value. If the Try contains an
    * exception it will be rethrown.
    *
    * @returns const reference to the contained value
    */
-  const T& value() const;
+  const T& value() const&;
 
   /*
    * If the Try contains an exception, rethrow it. Otherwise do nothing.
-- 
2.34.1