address comments

Snapchat · Jul 24, 2024 · 987cd72 · 987cd72
1 parent bf60220
commit 987cd72
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/support-lib/cpp/SharedFuture.hpp b/support-lib/cpp/SharedFuture.hpp
@@ -53,16 +53,24 @@ class SharedFuture {
         return await_resume();
     }
 
+    template <typename Func>
+    using ResultT = std::remove_cv_t<std::remove_reference_t<std::invoke_result_t<Func, const SharedFuture<T>&>>>;
+
     // Transform the result of this future into a new future. The behavior is same as Future::then except that
     // it doesn't consume the future, and can be called multiple times.
     template<typename Func>
-    SharedFuture<std::remove_cv_t<std::remove_reference_t<std::invoke_result_t<Func, const SharedFuture<T>&>>>> then(
-        Func transform) const {
+    Future<ResultT<Func>> then(Func transform) const {
         auto cpy = SharedFuture(*this); // retain copy during coroutine suspension
         co_await cpy.waitIgnoringExceptions();
         co_return transform(cpy);
     }
 
+    // Same as above but returns SharedFuture.
+    template<typename Func>
+    SharedFuture<ResultT<Func>> thenShared(Func transform) const {
+        return SharedFuture<ResultT<Func>>(then(std::move(transform)));
+    }
+
     // -- coroutine support implementation only; not intended externally --
 
     bool await_ready() const {

diff --git a/test-suite/handwritten-src/objc/tests/DBSharedFutureTest.mm b/test-suite/handwritten-src/objc/tests/DBSharedFutureTest.mm
@@ -37,28 +37,28 @@ - (void)testThen
     djinni::Promise<int> intPromise;
     djinni::SharedFuture<int> futureInt(intPromise.getFuture());
 
-    auto transformedInt = futureInt.then([](const auto& resolved) { return 2 * resolved.get(); });
+    auto transformedInt = futureInt.thenShared([](const auto& resolved) { return 2 * resolved.get(); });
 
     intPromise.setValue(42);
     XCTAssertEqual(transformedInt.get(), 84);
 
     // Also verify multiple consumers and chaining.
-    auto transformedString = futureInt.then([](const auto& resolved) { return std::to_string(resolved.get()); });
-    auto futurePlusOneTimesTwo = futureInt.then([](const auto& resolved) { return resolved.get() + 1; }).then([](const auto& resolved) {
+    auto transformedString = futureInt.thenShared([](const auto& resolved) { return std::to_string(resolved.get()); });
+    auto futurePlusOneTimesTwo = futureInt.then([](auto resolved) { return resolved.get() + 1; }).then([](auto resolved) {
         return 2 * resolved.get();
     });
-    auto futureStringLen = transformedString.then([](const auto& resolved) { return resolved.get().length(); }).toFuture();
+    auto futureStringLen = transformedString.then([](auto resolved) { return resolved.get().length(); });
 
     XCTAssertEqual(transformedString.get(), std::string("42"));
     XCTAssertEqual(futurePlusOneTimesTwo.get(), (42 + 1) * 2);
     XCTAssertEqual(futureStringLen.get(), 2);
 
     XCTAssertEqual(futureInt.get(), 42);
 
-    auto voidFuture = transformedString.then([](auto) {});
+    auto voidFuture = transformedString.thenShared([](auto) {});
     voidFuture.wait();
 
-    auto intFuture2 = voidFuture.then([](auto) { return 43; });
+    auto intFuture2 = voidFuture.thenShared([](auto) { return 43; });
     XCTAssertEqual(intFuture2.get(), 43);
 }
 
@@ -72,10 +72,10 @@ - (void)testException
 
     XCTAssertThrows(futureInt.get());
 
-    auto thenResult = futureInt.then([](const auto& resolved) { return resolved.get(); });
+    auto thenResult = futureInt.then([](auto resolved) { return resolved.get(); });
     XCTAssertThrows(thenResult.get());
 
-    auto withExceptionHandling = futureInt.then([](const auto& resolved) {
+    auto withExceptionHandling = futureInt.thenShared([](const auto& resolved) {
       try {
         return resolved.get();
       } catch (...) {