From 422a65cade2ae39b1eaed7a5afcee657787da19e Mon Sep 17 00:00:00 2001
From: q66 <daniel@octaforge.org>
Date: Wed, 22 Mar 2017 22:38:48 +0100
Subject: [PATCH] actually free coroutine stacks in dtor (more predictable and
 safer)

---
 ostd/coroutine.hh        |  6 +++++-
 ostd/internal/context.hh | 45 ++++++++++++++++++++++------------------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/ostd/coroutine.hh b/ostd/coroutine.hh
index 142ce73..ed68351 100644
--- a/ostd/coroutine.hh
+++ b/ostd/coroutine.hh
@@ -130,6 +130,7 @@ namespace detail {
     /* default case, yield returns args and takes a value */
     template<typename Y, typename R, typename ...A>
     struct coro_stor {
+        coro_stor(): p_func(nullptr) {}
         template<typename F>
         coro_stor(F &&func): p_func(std::forward<F>(func)) {}
 
@@ -171,6 +172,7 @@ namespace detail {
     /* yield takes a value but doesn't return any args */
     template<typename Y, typename R>
     struct coro_stor<Y, R> {
+        coro_stor(): p_func(nullptr) {}
         template<typename F>
         coro_stor(F &&func): p_func(std::forward<F>(func)) {}
 
@@ -199,6 +201,7 @@ namespace detail {
     /* yield doesn't take a value and returns args */
     template<typename Y, typename ...A>
     struct coro_stor<Y, void, A...> {
+        coro_stor(): p_func(nullptr) {}
         template<typename F>
         coro_stor(F &&func): p_func(std::forward<F>(func)) {}
 
@@ -232,6 +235,7 @@ namespace detail {
     /* yield doesn't take a value or return any args */
     template<typename Y>
     struct coro_stor<Y, void> {
+        coro_stor(): p_func(nullptr) {}
         template<typename F>
         coro_stor(F &&func): p_func(std::forward<F>(func)) {}
 
@@ -343,7 +347,7 @@ public:
     }
 
     template<typename SA = default_stack>
-    coroutine(std::nullptr_t, SA = SA{0}): base_t(), p_stor() {
+    coroutine(std::nullptr_t, SA = SA{}): base_t(), p_stor() {
         this->set_dead();
     }
 
diff --git a/ostd/internal/context.hh b/ostd/internal/context.hh
index 7e9130f..9070732 100644
--- a/ostd/internal/context.hh
+++ b/ostd/internal/context.hh
@@ -56,15 +56,17 @@ protected:
     /* coroutine context must be polymorphic */
     virtual ~coroutine_context() {
         unwind();
+        free_stack();
     }
 
     coroutine_context(coroutine_context const &) = delete;
     coroutine_context(coroutine_context &&c):
         p_stack(std::move(c.p_stack)), p_coro(c.p_coro), p_orig(c.p_orig),
-        p_except(std::move(c.p_except)), p_state(std::move(c.p_state))
+        p_except(std::move(c.p_except)), p_state(c.p_state), p_free(c.p_free)
     {
         c.p_coro = c.p_orig = nullptr;
         c.p_stack = { nullptr, 0 };
+        c.p_free = nullptr;
         c.set_dead();
     }
 
@@ -93,6 +95,11 @@ protected:
         p_orig = detail::ostd_jump_fcontext(p_orig, nullptr).ctx;
     }
 
+    void yield_done() {
+        set_dead();
+        p_orig = detail::ostd_jump_fcontext(p_orig, nullptr).ctx;
+    }
+
     bool is_hold() const {
         return (p_state == state::HOLD);
     }
@@ -122,6 +129,7 @@ protected:
         swap(p_orig, other.p_orig);
         swap(p_except, other.p_except);
         swap(p_state, other.p_state);
+        swap(p_free, other.p_free);
     }
 
     template<typename C, typename SA>
@@ -134,6 +142,7 @@ protected:
 
         p_coro = detail::ostd_make_fcontext(sp, asize, &context_call<C, SA>);
         new (sp) SA(std::move(sa));
+        p_free = &free_stack_call<SA>;
     }
 
 private:
@@ -171,11 +180,7 @@ private:
             return;
         }
         if (!p_orig) {
-            /* this coroutine never got to live :(
-             * let it call the entry point at least this once...
-             * this will kill the stack so we don't leak memory
-             */
-            coro_jump();
+            /* this coroutine never got to live :( */
             return;
         }
         detail::ostd_ontop_fcontext(
@@ -187,19 +192,18 @@ private:
     }
 
     template<typename SA>
-    void finish() {
-        set_dead();
-        detail::ostd_ontop_fcontext(
-            p_orig, this, [](detail::transfer_t t) -> detail::transfer_t {
-                auto &self = *(static_cast<coroutine_context *>(t.data));
-                auto &sa = *(static_cast<SA *>(self.get_stack_ptr<SA>()));
-                SA dsa{std::move(sa)};
-                /* in case it holds any state that needs destroying */
-                sa.~SA();
-                dsa.deallocate(self.p_stack);
-                return { nullptr, nullptr };
-            }
-        );
+    static void free_stack_call(void *data) {
+        auto &self = *(static_cast<coroutine_context *>(data));
+        auto &sa = *(static_cast<SA *>(self.get_stack_ptr<SA>()));
+        SA dsa{std::move(sa)};
+        sa.~SA();
+        dsa.deallocate(self.p_stack);
+    }
+
+    void free_stack() {
+        if (p_free) {
+            p_free(this);
+        }
     }
 
     template<typename C, typename SA>
@@ -223,7 +227,7 @@ private:
         }
         /* switch back, release stack */
 release:
-        self.template finish<SA>();
+        self.yield_done();
     }
 
     stack_context p_stack;
@@ -231,6 +235,7 @@ release:
     detail::fcontext_t p_orig = nullptr;
     std::exception_ptr p_except;
     state p_state = state::HOLD;
+    void (*p_free)(void *) = nullptr;
 };
 
 } /* namespace ostd */