From 18cfe5119f10ed8896358e1098a8c3a7ca1bd923 Mon Sep 17 00:00:00 2001
From: q66 <daniel@octaforge.org>
Date: Wed, 22 Mar 2017 18:23:29 +0100
Subject: [PATCH] track current task for coroutine scheduler separately

---
 build.sh                 |  2 +-
 ostd/concurrency.hh      | 59 ++++++++++++++++++++++------------------
 ostd/internal/context.hh | 16 ++++++++---
 src/concurrency.cc       | 14 ++++++++++
 4 files changed, 60 insertions(+), 31 deletions(-)
 create mode 100644 src/concurrency.cc

diff --git a/build.sh b/build.sh
index 8778e56..5232700 100755
--- a/build.sh
+++ b/build.sh
@@ -13,7 +13,7 @@ ASM_SOURCES="jump_all_gas make_all_gas ontop_all_gas"
 
 # c++ sources
 CXX_SOURCE_DIR="src"
-CXX_SOURCES="context_stack io"
+CXX_SOURCES="context_stack io concurrency"
 
 # output lib
 OSTD_LIB="libostd"
diff --git a/ostd/concurrency.hh b/ostd/concurrency.hh
index 8e1d809..160c1db 100644
--- a/ostd/concurrency.hh
+++ b/ostd/concurrency.hh
@@ -6,6 +6,8 @@
 #ifndef OSTD_CONCURRENCY_HH
 #define OSTD_CONCURRENCY_HH
 
+#include <vector>
+#include <list>
 #include <thread>
 #include <utility>
 #include <memory>
@@ -226,38 +228,38 @@ using simple_coroutine_scheduler =
 using protected_simple_coroutine_scheduler =
     basic_simple_coroutine_scheduler<stack_traits, true>;
 
-template<typename TR, bool Protected>
-struct basic_coroutine_scheduler {
-private:
-    struct task_cond;
-    struct task;
+namespace detail {
+    struct csched_task;
 
-    using tlist = std::list<task>;
-    using titer = typename tlist::iterator;
+    OSTD_EXPORT extern thread_local csched_task *current_csched_task;
 
-    struct task: coroutine_context {
+    struct csched_task: coroutine_context {
         std::function<void()> func;
-        task_cond *waiting_on = nullptr;
-        task *next_waiting = nullptr;
-        titer pos;
+        void *waiting_on = nullptr;
+        csched_task *next_waiting = nullptr;
+        typename std::list<csched_task>::iterator pos;
 
-        task() = delete;
-        task(task const &) = delete;
-        task(task &&) = delete;
-        task &operator=(task const &) = delete;
-        task &operator=(task &&) = delete;
+        csched_task() = delete;
+        csched_task(csched_task const &) = delete;
+        csched_task(csched_task &&) = delete;
+        csched_task &operator=(csched_task const &) = delete;
+        csched_task &operator=(csched_task &&) = delete;
 
         template<typename F, typename SA>
-        task(F &&f, SA &&sa): func(std::forward<F>(f)) {
+        csched_task(F &&f, SA &&sa): func(std::forward<F>(f)) {
             if (!func) {
                 this->set_dead();
                 return;
             }
-            this->make_context<task>(sa);
+            this->make_context<csched_task>(sa);
         }
 
         void operator()() {
-            this->call();
+            this->set_exec();
+            csched_task *curr = std::exchange(current_csched_task, this);
+            this->coro_jump();
+            current_csched_task = curr;
+            this->rethrow();
         }
 
         void yield() {
@@ -272,19 +274,24 @@ private:
             return this->is_dead();
         }
 
-        static task *current() {
-            auto ctx = coroutine_context::current();
-            task *t = dynamic_cast<task *>(ctx);
-            if (!t) {
-                std::terminate();
-            }
-            return t;
+        static csched_task *current() {
+            return current_csched_task;
         }
 
         void resume_call() {
             func();
         }
     };
+}
+
+template<typename TR, bool Protected>
+struct basic_coroutine_scheduler {
+private:
+    struct task_cond;
+    using task = detail::csched_task;
+
+    using tlist = std::list<task>;
+    using titer = typename tlist::iterator;
 
     struct task_cond {
         task_cond() = delete;
diff --git a/ostd/internal/context.hh b/ostd/internal/context.hh
index eb21be3..7e9130f 100644
--- a/ostd/internal/context.hh
+++ b/ostd/internal/context.hh
@@ -75,15 +75,13 @@ protected:
     }
 
     void call() {
-        p_state = state::EXEC;
+        set_exec();
         /* switch to new coroutine */
         coroutine_context *curr = std::exchange(detail::coro_current, this);
         coro_jump();
         /* switch back to previous */
         detail::coro_current = curr;
-        if (p_except) {
-            std::rethrow_exception(std::move(p_except));
-        }
+        rethrow();
     }
 
     void coro_jump() {
@@ -107,6 +105,16 @@ protected:
         p_state = state::TERM;
     }
 
+    void set_exec() {
+        p_state = state::EXEC;
+    }
+
+    void rethrow() {
+        if (p_except) {
+            std::rethrow_exception(std::move(p_except));
+        }
+    }
+
     void swap(coroutine_context &other) noexcept {
         using std::swap;
         swap(p_stack, other.p_stack);
diff --git a/src/concurrency.cc b/src/concurrency.cc
new file mode 100644
index 0000000..35f7678
--- /dev/null
+++ b/src/concurrency.cc
@@ -0,0 +1,14 @@
+/* Concurrency C implementation bits.
+ *
+ * This file is part of OctaSTD. See COPYING.md for futher information.
+ */
+
+#include "ostd/concurrency.hh"
+
+namespace ostd {
+namespace detail {
+
+OSTD_EXPORT thread_local csched_task *current_csched_task = nullptr;
+
+} /* namespace detail */
+} /* namespace ostd */