stackfull-coroutines.cpp

#include <iostream>
#include <csetjmp>

class CoroutineImpl {
public:
    using CalledFunction = void(*)(CoroutineImpl*);

    CoroutineImpl(CalledFunction func) {
        m_stackSize = 8 * 1024;
        m_stackPointer = malloc(m_stackSize);
        m_func = func;
    }

    ~CoroutineImpl() {
        free(m_stackPointer);
    }

    void start() {
        if(!m_ended) return;

        m_ended = false;
        if(setjmp(m_calleeContext) == 0) {
            // dirty hack, the this pointer will be
            // invalidated when I move the %rps, so
            // I store it in a register.
            asm("movq %0,%%r10" : : "r"(this));
            asm("movq %0,%%rsp\n"
                "movq %%rsp,%%rbp" : : "r"(((char*)m_stackPointer) + m_stackSize));
            {
                // then I use another function to move
                // it back in the stack, so I have no
                // issue if for some reason I write to
                // that register.

                CoroutineImpl* self = nullptr;
                asm volatile("movq %%r10,%0" : "=r"(self));
                coroutineStart(self);
            }
        }
    }
    void resume() {
        if(m_ended) return;

        if(setjmp(m_calleeContext) == 0) {
            longjmp(m_calledContext, 1);
        }
    }
    void restore() {
        if(m_ended) return;

        if(setjmp(m_calledContext) == 0) {
            longjmp(m_calleeContext, 1);
        }
    }
    bool ended() {
        return m_ended;
    }

private:
    static void coroutineStart(CoroutineImpl* self) {
        self->m_func(self);

        self->m_ended = true;
        longjmp(self->m_calleeContext, 1);
    }

    bool m_ended = true;

    unsigned m_stackSize;
    void* m_stackPointer;
    CalledFunction m_func;

    jmp_buf m_calledContext;
    jmp_buf m_calleeContext;
};

template <class Lambda, class Ret>
class Generator : protected CoroutineImpl {
public:
    Generator(Lambda lambda, bool lazy) :
        CoroutineImpl(Generator::coroStart),
        m_lambda(lambda),
        m_lazy(lazy) {

        start();
    }

    Generator(Lambda lambda) : Generator<Lambda, Ret>(lambda, false) {}

    Generator(const Generator<Lambda, Ret>&) = delete;
    Generator(Generator<Lambda, Ret>&&) = default;

    bool hasNext() {
        if(ended()) return false;

        if(!m_next) {
            resume(); // will block until restore is called

            return m_next;
        }
        return true;
    }
    Ret& next() {
        if(hasNext()) {
            m_next = false;
            return m_ret;
        }
        throw "NoMoreStuff";
    }
    void yield(Ret ret) {
        m_next = true;
        m_ret = ret;

        restore(); // will block until resume is called
    }

protected:
    static void coroStart(CoroutineImpl* self) {
        static_cast<Generator<Lambda, Ret>*>(self)->coroStart();
    }

    void coroStart() {
        if(m_lazy) restore(); // will block until resume is called
        m_lambda(*this);
    }

private:
    Lambda m_lambda;
    bool m_lazy;

    bool m_next = false;
    Ret m_ret;
};

template <class Ret, class Lambda>
Generator<Lambda, Ret> G(Lambda l) {
    return Generator<Lambda, Ret>(l);
}

int main() {
    auto generator = G<int>([](auto& self) {
        for(int i = 0; i < 100; i++) {
            if(i % 12 == 0) {
                self.yield(i);
            }
        }
    });
    while(generator.hasNext()) {
        std::cout << "Gen returned: " << generator.next() << std::endl;
    }
    return 0;
}