From 3597c587da97b02466811339af1e9ced2e73d988 Mon Sep 17 00:00:00 2001
From: Dylan Yudaken <dylany@fb.com>
Date: Fri, 15 Sep 2017 10:13:16 -0700
Subject: [PATCH] Allow awaiting on a folly::Optional that returns a move-only
 type

Summary:
await_resume is only called once, so this allows it to move the value out.
At the same time remove a redundant clear (but keep the existing requirement that the promise type is an OptionalPromise), and clean up the tests.
Also add a test to make sure the coroutine is cleaned up

Reviewed By: ericniebler

Differential Revision: D5834861

fbshipit-source-id: 7ad487e818969cdf6fe27c9e82931aa247daf4e4
---
 folly/Optional.h                      | 15 ++++++------
 folly/test/OptionalCoroutinesTest.cpp | 35 +++++++++++++++++++++------
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/folly/Optional.h b/folly/Optional.h
index 2693423d..2bb2d658 100644
--- a/folly/Optional.h
+++ b/folly/Optional.h
@@ -598,13 +598,14 @@ struct OptionalAwaitable {
     return o_.hasValue();
   }
   Value await_resume() {
-    return o_.value();
+    return std::move(o_.value());
   }
-  template <typename CoroHandle>
-  void await_suspend(CoroHandle h) const {
-    // make sure the coroutine returns an empty Optional:
-    h.promise().value_->clear();
-    // Abort the rest of the coroutine:
+
+  // Explicitly only allow suspension into an OptionalPromise
+  template <typename U>
+  void await_suspend(
+      std::experimental::coroutine_handle<OptionalPromise<U>> h) const {
+    // Abort the rest of the coroutine. resume() is not going to be called
     h.destroy();
   }
 };
@@ -617,7 +618,7 @@ detail::OptionalAwaitable<Value>
 }
 } // namespace folly
 
-// This makes std::optional<Value> useable as a coroutine return type..
+// This makes folly::Optional<Value> useable as a coroutine return type..
 FOLLY_NAMESPACE_STD_BEGIN
 namespace experimental {
 template <typename Value, typename... Args>
diff --git a/folly/test/OptionalCoroutinesTest.cpp b/folly/test/OptionalCoroutinesTest.cpp
index 92c6ba60..9d77718a 100644
--- a/folly/test/OptionalCoroutinesTest.cpp
+++ b/folly/test/OptionalCoroutinesTest.cpp
@@ -16,6 +16,7 @@
 
 #include <folly/Optional.h>
 #include <folly/Portability.h>
+#include <folly/ScopeGuard.h>
 #include <folly/portability/GTest.h>
 
 #if FOLLY_HAS_COROUTINES
@@ -27,8 +28,10 @@ Optional<int> f1() {
 Optional<double> f2(int x) {
   return 2.0 * x;
 }
-Optional<int> f3(int x, double y) {
-  return (int)(x + y);
+
+// move-only type
+Optional<std::unique_ptr<int>> f3(int x, double y) {
+  return std::make_unique<int>((int)(x + y));
 }
 
 TEST(Optional, CoroutineSuccess) {
@@ -38,8 +41,8 @@ TEST(Optional, CoroutineSuccess) {
     auto y = co_await f2(x);
     EXPECT_EQ(2.0 * 7, y);
     auto z = co_await f3(x, y);
-    EXPECT_EQ((int)(2.0 * 7 + 7), z);
-    co_return z;
+    EXPECT_EQ((int)(2.0 * 7 + 7), *z);
+    co_return* z;
   }();
   EXPECT_TRUE(r0.hasValue());
   EXPECT_EQ(21, *r0);
@@ -54,7 +57,7 @@ TEST(Optional, CoroutineFailure) {
     auto x = co_await f1();
     auto y = co_await f2(x);
     auto z = co_await f4(x, y);
-    EXPECT_FALSE(true);
+    ADD_FAILURE();
     co_return z;
   }();
   EXPECT_TRUE(!r1.hasValue());
@@ -68,14 +71,30 @@ TEST(Optional, CoroutineException) {
   try {
     auto r2 = []() -> Optional<int> {
       auto x = co_await throws();
-      EXPECT_FALSE(true);
+      ADD_FAILURE();
       co_return x;
     }();
-    EXPECT_FALSE(true);
+    ADD_FAILURE();
   } catch (/* nolint */ int i) {
     EXPECT_EQ(42, i);
   } catch (...) {
-    EXPECT_FALSE(true);
+    ADD_FAILURE();
   }
 }
+
+// this test makes sure that the coroutine is destroyed properly
+TEST(Optional, CoroutineCleanedUp) {
+  int count_dest = 0;
+  auto r = [&]() -> Optional<int> {
+    SCOPE_EXIT {
+      ++count_dest;
+    };
+    auto x = co_await folly::Optional<int>();
+    ADD_FAILURE() << "Should not be resuming";
+    co_return x;
+  }();
+  EXPECT_FALSE(r.hasValue());
+  EXPECT_EQ(1, count_dest);
+}
+
 #endif
-- 
2.34.1