Fix coroutine continuation cleanup

support-lib/cpp/Future.hpp

@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include "expected.hpp"
+
 #include <atomic>
 #include <functional>
 #include <memory>
@@ -364,34 +366,63 @@ 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{
-        template <typename V>
+
+    struct PromiseType: PromiseTypeBase<PromiseType>{
+        template <typename V, typename = std::enable_if_t<std::is_convertible_v<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

support-lib/cpp/expected.hpp

@@ -17,6 +17,7 @@ namespace djinni {
 
 using ::std::expected;
 using ::std::unexpected;
+inline constexpr ::std::unexpect_t unexpect{};
 
 }
 
@@ -28,6 +29,7 @@ namespace djinni {
 
 using ::tl::unexpected;
 using ::tl::expected;
+inline constexpr ::tl::unexpect_t unexpect{};
 
 }
 

test-suite/handwritten-src/objc/tests/DBCppFutureTests.mm

@@ -0,0 +1,70 @@
+#import <Foundation/Foundation.h>
+#import <XCTest/XCTest.h>
+
+#include <Future.hpp>
+
+#ifdef DJINNI_FUTURE_HAS_COROUTINE_SUPPORT
+
+namespace {
+
+template<typename Functor>
+struct OnCleanup {
+    OnCleanup(Functor&& functor)
+    :functor{std::move(functor)}
+    {}
+    ~OnCleanup() {
+        functor();
+    }
+    Functor functor;
+};
+
+djinni::Future<void> inner_coroutine(std::vector<int>& cleanup_ids, int id, djinni::Future<void> suspendOn) {
+    OnCleanup cleanup{
+        [&] {
+            cleanup_ids.push_back(id);
+        }
+    };
+
+    co_await suspendOn;
+    co_return; // do not remove!
+}
+
+djinni::Future<void> outer_coroutine(std::vector<int>& cleanup_ids, int id, djinni::Future<void> suspendOn) {
+    OnCleanup cleanup{
+        [&] {
+            cleanup_ids.push_back(id);
+        }
+    };
+
+    co_await inner_coroutine(cleanup_ids, id + 1, std::move(suspendOn));
+    co_return; // do not remove!
+}
+
+}
+
+#endif
+
+@interface DBCppFutureTests : XCTestCase
+@end
+
+@implementation DBCppFutureTests
+
+#ifdef DJINNI_FUTURE_HAS_COROUTINE_SUPPORT
+- (void) testFutureCoroutines_cleanupOrder {
+    std::vector<int> cleanupIds{};
+
+    djinni::Promise<void> djinniPromise{};
+
+    auto coroutineFuture = outer_coroutine(cleanupIds, 0, djinniPromise.getFuture());
+    XCTAssertFalse(coroutineFuture.isReady());
+
+    djinniPromise.setValue();
+    XCTAssertTrue(coroutineFuture.isReady());
+
+    XCTAssertEqual(cleanupIds.size(), 2);
+    XCTAssertEqual(cleanupIds[0], 1); // the inner coroutine should be cleaned up first
+    XCTAssertEqual(cleanupIds[1], 0); // ... then the outer
+}
+#endif
+
+@end