Playing with generators in C++

August 18, 2019 - lutz

I wanted to play with generators and coroutines in C++. Coroutines are a feature I really like in python. In a nutshell a coroutine is a function that can suspend itself and return a value on each suspend. This is called yielding. The next iteration of C++ will include coroutines and add the keyword co_yield to perform yielding. In this post I’ll show my implementation of a coroutine-like generator.

Stack branching

A generator function needs to be suspended and continued. To achieve this the state of the generator function has to be saved when yield is called and the yielded values needs to be saved in a fashion that it can be pushed to the calling context. However, since the calling context can call functions after the generator has yielded (thus overwriting the stack) the generator’s stack must be saved somewhere else. In practice this means that before executing the generator a different stack has to be used. When calling a regular function or a generator a different stack has to be utilized within the callee; thus the term stack brancing. To get this working I’ve utilized UNIX’ ucontext. ucontext was originally designed to enable user mode context switching to implement multitasking. For modern day implementations it might not be the best choice for context switching as it is kinda slow (see). Also the generator function is wrapped inside a Generator object that implements begin() and end(). The forward-iterator returned from begin() runs the generator function when operator++() is called (and once when constructed). The generator function is passed a typed yielder that handles the context switching from the generator function to the caller. The rest is pretty straight forward.

#include <iostream>
#include "Generator.h"

int main() {
    auto gen = [](auto& t) {
        for (int i{0}; i < 10; ++i) {
            t.yield(i);
        }
    };
    for (auto i : Generator<int>{gen}) {
        std::cout << i << "\n";
    }
    for (auto i : Generator<int>{gen}) {
        std::cout << i << "\n";
        break;
    }
    std::cout << "done" << std::endl;
    return 0;
}

Generator.h

#include <ucontext.h>

#include <iterator>
#include <optional>
#include <functional>
#include <memory>
#include <future>
#include <exception>
#include <algorithm>

template<typename T>
struct Yielder {
protected:
    std::optional<T> val;
    ucontext_t& caller_context;
    ucontext_t& callee_context;
    std::exception_ptr eptr{};

    Yielder(ucontext_t& _caller_context, ucontext_t& _callee_context)
        : caller_context{_caller_context}
        , callee_context{_callee_context}
    {}

public:
    template<typename... Args>
    void yield(Args &&... args) {
        val.emplace(std::forward<Args>(args)...);
        swapcontext(&callee_context, &caller_context);
        if (eptr) {
            std::rethrow_exception(eptr);
        }
    }

    auto operator*() -> T& {
        return val.operator*();
    }

    auto operator->() -> T& {
        return val.operator->();
    }

    auto operator*() const -> T const& {
        return val.operator*();
    }

    auto operator->() const -> T const* {
        return val.operator->();
    }
};

template<typename T, int stack_size=4096>
struct Generator {
private:
    std::function<void(Yielder<T>&)> func;
public:
    struct end_iterator {};
    struct iterator : Yielder<T>, std::iterator<std::forward_iterator_tag, T> {
        using super = std::iterator<std::forward_iterator_tag, T>;
        using value_type = typename super::value_type;
        using reference = typename super::reference;
        using pointer = typename super::reference;

    private:
        struct CleanupStack : std::exception {};

        std::function<void(Yielder<T>&)> func;
        std::array<std::byte, stack_size> stack;

        ucontext_t caller_context;
        ucontext_t callee_context;

        bool returned {false};
        static void call(iterator* iter) noexcept {
            try {
                iter->func(*iter);
                iter->returned = true;
            } catch (CleanupStack const&) {}
            iter->val.reset();
        }
    public:
        iterator(std::function<void(Yielder<T>&)>&& f) : Yielder<T>{caller_context, callee_context}, func{std::move(f)} {
            getcontext(&caller_context);
            getcontext(&callee_context);
            callee_context.uc_stack.ss_sp   = stack.data();
            callee_context.uc_stack.ss_size = stack.size();
            callee_context.uc_link = &caller_context;
            makecontext(&callee_context, reinterpret_cast<void(*)()>(call), 1, this);
            ++(*this);
        }

        ~iterator() {
            if (not returned) {
                // in this case the generator context has to be torn down because it cannot be continued to its end
                // setting the eptr will rethrow the exception the next time just before the function is continued thus tearing down the callees stack
                try {
                    throw CleanupStack{};
                } catch (...) {
                    this->eptr = std::current_exception();                                
                }
                ++(*this);
            }
        }

        iterator& operator++() {
            swapcontext(&caller_context, &callee_context);
            return *this;
        }

        template<typename Other>
        bool operator!=(Other const&) const {
            return this->val.has_value();
        }
    };

    template<typename Functor>
    Generator(Functor&& f) : func{std::forward<Functor>(f)} {}

    iterator begin() {
        return iterator(std::move(func));
    }
    end_iterator end() {
        return {};
    }
};

Stack Maintenance

Loops can be left early with a break statement thus there is a necessity to perform cleanup of the generator’s stack. In my code this is achieved by throwing an exception into the generator before the completion of the call to yield. The Exception is then caught internally and there is no need for the caller to worry about anything.