From fb2a27a30f30b5d612ea4728f9216b127ba0145d Mon Sep 17 00:00:00 2001
From: Lee Howes <lwh@fb.com>
Date: Tue, 26 Dec 2017 13:34:26 -0800
Subject: [PATCH] Ensure that returning a semifuture from a continuation works
 correctly.

Summary: Returning a SemiFuture from a continuation should work by correctly checking the types and returning a folly::Future on the same executor as the original future that .then was applied to.

Reviewed By: yfeldblum

Differential Revision: D6597273

fbshipit-source-id: cf2016a344d4b29f1d31c1da20c89df5b4cfe64e
---
 folly/futures/Future-pre.h            | 28 ++++++++++++++++----
 folly/futures/test/SemiFutureTest.cpp | 38 +++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/folly/futures/Future-pre.h b/folly/futures/Future-pre.h
index 82a8684c..560b4df6 100644
--- a/folly/futures/Future-pre.h
+++ b/folly/futures/Future-pre.h
@@ -45,6 +45,24 @@ struct isFuture<Future<T>> : std::true_type {
   typedef T Inner;
 };
 
+template <typename T>
+struct isFutureOrSemiFuture : std::false_type {
+  using Inner = typename Unit::Lift<T>::type;
+  using Return = Inner;
+};
+
+template <typename T>
+struct isFutureOrSemiFuture<Future<T>> : std::true_type {
+  typedef T Inner;
+  using Return = Future<Inner>;
+};
+
+template <typename T>
+struct isFutureOrSemiFuture<SemiFuture<T>> : std::true_type {
+  typedef T Inner;
+  using Return = SemiFuture<Inner>;
+};
+
 template <typename T>
 struct isTry : std::false_type {};
 
@@ -109,7 +127,7 @@ struct callableResult {
           callableWith<F, Try<T>&&>::value,
           detail::argResult<true, F, Try<T>&&>,
           detail::argResult<true, F, Try<T>&>>::type>::type>::type>::type Arg;
-  typedef isFuture<typename Arg::Result> ReturnsFuture;
+  typedef isFutureOrSemiFuture<typename Arg::Result> ReturnsFuture;
   typedef Future<typename ReturnsFuture::Inner> Return;
 };
 
@@ -118,7 +136,7 @@ struct Extract : Extract<decltype(&L::operator())> { };
 
 template <typename Class, typename R, typename... Args>
 struct Extract<R(Class::*)(Args...) const> {
-  typedef isFuture<R> ReturnsFuture;
+  typedef isFutureOrSemiFuture<R> ReturnsFuture;
   typedef Future<typename ReturnsFuture::Inner> Return;
   typedef typename ReturnsFuture::Inner RawReturn;
   typedef typename ArgType<Args...>::FirstArg FirstArg;
@@ -126,7 +144,7 @@ struct Extract<R(Class::*)(Args...) const> {
 
 template <typename Class, typename R, typename... Args>
 struct Extract<R(Class::*)(Args...)> {
-  typedef isFuture<R> ReturnsFuture;
+  typedef isFutureOrSemiFuture<R> ReturnsFuture;
   typedef Future<typename ReturnsFuture::Inner> Return;
   typedef typename ReturnsFuture::Inner RawReturn;
   typedef typename ArgType<Args...>::FirstArg FirstArg;
@@ -134,7 +152,7 @@ struct Extract<R(Class::*)(Args...)> {
 
 template <typename R, typename... Args>
 struct Extract<R (*)(Args...)> {
-  typedef isFuture<R> ReturnsFuture;
+  typedef isFutureOrSemiFuture<R> ReturnsFuture;
   typedef Future<typename ReturnsFuture::Inner> Return;
   typedef typename ReturnsFuture::Inner RawReturn;
   typedef typename ArgType<Args...>::FirstArg FirstArg;
@@ -142,7 +160,7 @@ struct Extract<R (*)(Args...)> {
 
 template <typename R, typename... Args>
 struct Extract<R (&)(Args...)> {
-  typedef isFuture<R> ReturnsFuture;
+  typedef isFutureOrSemiFuture<R> ReturnsFuture;
   typedef Future<typename ReturnsFuture::Inner> Return;
   typedef typename ReturnsFuture::Inner RawReturn;
   typedef typename ArgType<Args...>::FirstArg FirstArg;
diff --git a/folly/futures/test/SemiFutureTest.cpp b/folly/futures/test/SemiFutureTest.cpp
index b5c059da..c51a6fe9 100644
--- a/folly/futures/test/SemiFutureTest.cpp
+++ b/folly/futures/test/SemiFutureTest.cpp
@@ -202,6 +202,44 @@ TEST(SemiFuture, MakeFutureFromSemiFuture) {
   ASSERT_EQ(result, 42);
 }
 
+TEST(SemiFuture, MakeFutureFromSemiFutureReturnFuture) {
+  folly::EventBase e;
+  Promise<int> p;
+  int result{0};
+  auto f = p.getSemiFuture();
+  auto future = std::move(f).via(&e).then([&](int value) {
+    result = value;
+    return folly::makeFuture(std::move(value));
+  });
+  e.loop();
+  EXPECT_EQ(result, 0);
+  EXPECT_FALSE(future.isReady());
+  p.setValue(42);
+  e.loop();
+  EXPECT_TRUE(future.isReady());
+  ASSERT_EQ(future.value(), 42);
+  ASSERT_EQ(result, 42);
+}
+
+TEST(SemiFuture, MakeFutureFromSemiFutureReturnSemiFuture) {
+  folly::EventBase e;
+  Promise<int> p;
+  int result{0};
+  auto f = p.getSemiFuture();
+  auto future = std::move(f).via(&e).then([&](int value) {
+    result = value;
+    return folly::makeSemiFuture(std::move(value));
+  });
+  e.loop();
+  EXPECT_EQ(result, 0);
+  EXPECT_FALSE(future.isReady());
+  p.setValue(42);
+  e.loop();
+  EXPECT_TRUE(future.isReady());
+  ASSERT_EQ(future.value(), 42);
+  ASSERT_EQ(result, 42);
+}
+
 TEST(SemiFuture, MakeFutureFromSemiFutureLValue) {
   folly::EventBase e;
   Promise<int> p;
-- 
2.34.1