Fix continuations in djinni::Future coroutines

Snapchat · Jul 15, 2024 · 5793459 · 5793459
1 parent fe7e06e
commit 5793459
Showing 1 changed file with 41 additions and 9 deletions.
diff --git a/support-lib/cpp/Future.hpp b/support-lib/cpp/Future.hpp
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include "expected.hpp"
+
 #include <atomic>
 #include <functional>
 #include <memory>
@@ -364,34 +366,64 @@ class Future {
         return true;
     }
 
+    template<typename ConcretePromise>
     struct PromiseTypeBase {
         Promise<T> _promise;
+        std::optional<djinni::expected<T, std::exception_ptr>> _result{};
+
+        struct FinalAwaiter {
+            constexpr bool await_ready() const noexcept {
+                return false;
+            }
+            bool await_suspend(std::coroutine_handle<ConcretePromise> finished) const noexcept {
+                auto& promise_type = finished.promise();
+                if (*promise_type._result) {
+                    if constexpr (std::is_void_v<T>) {
+                        promise_type._promise.setValue();
+                    } else {
+                        promise_type._promise.setValue(std::move(**promise_type._result));
+                    }
+                } else {
+                    promise_type._promise.setException(std::move(promise_type._result->error()));
+                }
+                return false;
+            }
+            constexpr void await_resume() const noexcept {}
+        };
 
-        detail::SuspendNever initial_suspend() { return {}; }
-        detail::SuspendNever final_suspend() noexcept { return {}; }
+        constexpr detail::SuspendNever initial_suspend() const noexcept { return {}; }
+        FinalAwaiter final_suspend() const noexcept { return {}; }
 
         Future<T> get_return_object() noexcept {
             return _promise.getFuture();
         }
         void unhandled_exception() {
-            _promise.setException(std::current_exception());
+            _result.emplace(djinni::unexpect, std::current_exception());
         }
     };
-    template <typename U>
-    struct PromiseType: PromiseTypeBase{
+
+    struct PromiseType: PromiseTypeBase<PromiseType>{
         template <typename V>
+        requires(std::convertible_to<V, T>)
         void return_value(V&& value) {
-            this->_promise.setValue(std::forward<V>(value));
+            this->_result.emplace(std::forward<V>(value));
+        }
+        void return_value(T&& value) {
+            this->_result.emplace(std::move(value));
+        }
+        void return_value(const T& value) {
+            this->_result.emplace(value);
         }
     };
-    using promise_type = PromiseType<T>;
+    using promise_type = PromiseType;
 #endif
 };
 
 #if defined(DJINNI_FUTURE_HAS_COROUTINE_SUPPORT)
-template<> template<> struct Future<void>::PromiseType<void>:PromiseTypeBase {
+template<>
+struct Future<void>::PromiseType : PromiseTypeBase<PromiseType> {
     void return_void() {
-        this->_promise.setValue();
+        _result.emplace();
     }
 };
 #endif