Common/Fiber: Implement Rewind on Boost Context.

2 files changed, 39 insertions, 2 deletions

diff --git a/src/common/fiber.cpp b/src/common/fiber.cpp
index f61479e13..6ea314d75 100644
--- a/src/common/fiber.cpp
+++ b/src/common/fiber.cpp
@@ -113,7 +113,11 @@ std::shared_ptr<Fiber> Fiber::ThreadToFiber() {
 
 struct Fiber::FiberImpl {
     alignas(64) std::array<u8, default_stack_size> stack;
+    u8* stack_limit;
+    alignas(64) std::array<u8, default_stack_size> rewind_stack;
+    u8* rewind_stack_limit;
     boost::context::detail::fcontext_t context;
+    boost::context::detail::fcontext_t rewind_context;
 };
 
 void Fiber::start(boost::context::detail::transfer_t& transfer) {
@@ -125,21 +129,43 @@ void Fiber::start(boost::context::detail::transfer_t& transfer) {
     UNREACHABLE();
 }
 
+void Fiber::onRewind(boost::context::detail::transfer_t& [[maybe_unused]] transfer) {
+    ASSERT(impl->context != nullptr);
+    impl->context = impl->rewind_context;
+    impl->rewind_context = nullptr;
+    u8* tmp = impl->stack_limit;
+    impl->stack_limit = impl->rewind_stack_limit;
+    impl->rewind_stack_limit = tmp;
+    rewind_point(rewind_parameter);
+    UNREACHABLE();
+}
+
 void Fiber::FiberStartFunc(boost::context::detail::transfer_t transfer) {
     auto fiber = static_cast<Fiber*>(transfer.data);
     fiber->start(transfer);
 }
 
+void Fiber::RewindStartFunc(boost::context::detail::transfer_t transfer) {
+    auto fiber = static_cast<Fiber*>(transfer.data);
+    fiber->onRewind(transfer);
+}
+
 Fiber::Fiber(std::function<void(void*)>&& entry_point_func, void* start_parameter)
     : guard{}, entry_point{std::move(entry_point_func)}, start_parameter{start_parameter},
       previous_fiber{} {
     impl = std::make_unique<FiberImpl>();
-    u8* stack_limit = impl->stack.data();
-    u8* stack_base = stack_limit + default_stack_size;
+    impl->stack_limit = impl->stack.data();
+    impl->rewind_stack_limit = impl->rewind_stack.data();
+    u8* stack_base = impl->stack_limit + default_stack_size;
     impl->context =
         boost::context::detail::make_fcontext(stack_base, impl->stack.size(), FiberStartFunc);
 }
 
+void Fiber::SetRewindPoint(std::function<void(void*)>&& rewind_func, void* start_parameter) {
+    rewind_point = std::move(rewind_func);
+    rewind_parameter = start_parameter;
+}
+
 Fiber::Fiber() {
     impl = std::make_unique<FiberImpl>();
 }
@@ -161,6 +187,15 @@ void Fiber::Exit() {
     guard.unlock();
 }
 
+void Fiber::Rewind() {
+    ASSERT(rewind_point);
+    ASSERT(impl->rewind_context == nullptr);
+    u8* stack_base = impl->rewind_stack_limit + default_stack_size;
+    impl->rewind_context =
+        boost::context::detail::make_fcontext(stack_base, impl->stack.size(), RewindStartFunc);
+    boost::context::detail::jump_fcontext(impl->rewind_context, this);
+}
+
 void Fiber::YieldTo(std::shared_ptr<Fiber> from, std::shared_ptr<Fiber> to) {
     ASSERT_MSG(from != nullptr, "Yielding fiber is null!");
     ASSERT_MSG(to != nullptr, "Next fiber is null!");
diff --git a/src/common/fiber.h b/src/common/fiber.h
index 3bbd506b5..cab7bc4b5 100644
--- a/src/common/fiber.h
+++ b/src/common/fiber.h
@@ -70,8 +70,10 @@ private:
     static void FiberStartFunc(void* fiber_parameter);
     static void RewindStartFunc(void* fiber_parameter);
 #else
+    void onRewind(boost::context::detail::transfer_t& transfer);
     void start(boost::context::detail::transfer_t& transfer);
     static void FiberStartFunc(boost::context::detail::transfer_t transfer);
+    static void RewindStartFunc(boost::context::detail::transfer_t transfer);
 #endif
 
     struct FiberImpl;