Cooperative concurrency

#include <cassert>
#include <functional>
#include <iostream>
#include <vector>

struct ISpringBoard {
    [[noreturn]] virtual void Run() = 0;
};

struct Context;

extern "C" {
    void SwitchContext(Context* from, Context to);
    void SetupContext(Context* self, void* stackTop, void* springBoard,
                      void (*runSpringBoard)(void*, void*, void*, void*, void*, void*, void*));
}

asm(R"(
SwitchContext:
    # Callee-saved registers according to x86_64 ABI
    # rbx, rsp, rbp, r12, r13, r14, r15
    pushq %rbx
    pushq %rbp
    pushq %r12
    pushq %r13
    pushq %r14
    pushq %r15

    movq %rsp, (%rdi)
    movq %rsi, %rsp

    popq %r15
    popq %r14
    popq %r13
    popq %r12
    popq %rbp
    popq %rbx

    ret

SetupContext:
    movq %rsp, %rax;
    movq %rsi, %rsp;

    pushq $0; # Stack alignment
    pushq %rdx;
    pushq $0; # Fake return address
    pushq %rcx;

    pushq $0; # rbx
    pushq $0; # rbp
    pushq $0; # r12
    pushq $0; # r13
    pushq $0; # r14
    pushq $0; # r15

    movq %rsp, (%rdi);

    movq %rax, %rsp;
    ret
)");


struct Context {
    void* rsp;

    void Setup(ISpringBoard& springBoard, void* stackTop) {
        SetupContext(this, stackTop, &springBoard, &Context::RunSpringBoard);
    }

    void SwitchTo(Context to) {
        SwitchContext(this, to);
    }

private:
    [[noreturn]] static void RunSpringBoard(void*, void*, void*, void*, void*, void*, void* springBoardRaw) {
        auto springBoard = static_cast<ISpringBoard*>(springBoardRaw);
        springBoard->Run();
    }
};

class Coroutine: private ISpringBoard {
public:
    Coroutine(std::function<void()> runnable, size_t stackSize)
        : runnable(std::move(runnable))
        , stack(stackSize) {
        context.Setup(*this, stack.data() + stackSize);
    }

    void Resume() {
        auto prev = std::exchange(currentCoroutine, this);
        context.SwitchTo(context);
        currentCoroutine = prev;
    }

    static void Suspend() {
        assert(currentCoroutine);
        auto& currentContext = currentCoroutine->context;
        currentContext.SwitchTo(currentContext);
    }

private:
    [[noreturn]] void Run() override {
        runnable();
        context.SwitchTo(context);
        std::abort(); // Finished corouting resumed
    }

    std::function<void()> runnable;
    Context context;
    std::vector<char> stack;

    static thread_local Coroutine *currentCoroutine;
};

thread_local Coroutine *Coroutine::currentCoroutine = nullptr;

int main() {
    Coroutine a([] {
        std::cout << "2" << std::endl;
        Coroutine::Suspend();
        std::cout << "6" << std::endl;
    }, 50 << 10);

    Coroutine b([] {
        std::cout << "3" << std::endl;
        Coroutine::Suspend();
        std::cout << "5" << std::endl;
    }, 20 << 10);

    std::cout << "1" << std::endl;
    a.Resume();
    b.Resume();
    std::cout << "4" << std::endl;
    b.Resume();
    a.Resume();
}