Coroutines - A Simple Generator

Before I demonstrated a simple coroutine. Here I want to move on to a slightly more realworld usecase. You will see how to create a simple generator.

Let say I want to write a program that will generate a simple list of prime numbers. To make things simpler lets say we have the function:

bool isPrime(uint64_t num)
{
    if (num < 2) {
        return false;
    }
    if (num == 2) {
        return true;
    }
    if (num % 2 == 0) {
        return false;
    }
    auto topSearch  = static_cast<uint64_t>(sqrt(num));
    for (auto i = 3ul; i <= topSearch; i+=2ul) {
        if ((num % i) == 0) {
            return false;
        }
    }
    return true;
}

So given a number, we can easily check if it is prime or not.

The Naïve Case

Now we can use this function to generate the first 100 primes using the following code.

//printf out the first n prime numbers
void nPrimes(size_t n)
{
    uint64_t num = 0;
    for (size_t i = 0; i < n; ++i, ++num) {
        while (!isPrime(num)) {
            ++num;
        }
        std::cout << num << " is prime\n";
    }
}

int main()
{
    nPrimes(100);
    return 0;
}

This will as required print out the first 100 primes. It is in a nice sequential easy to understand format. There is only one problem. The primes will be all printed to the standard output. We might want to store them for future calculations or some other purpose. nPrimes does not allow that.

now we could change the function to something like:

void nPrimes(size_t n, std::function<void(uint64_t prime));

And using the callback we can do what we want with our newly minted prime number. However once started we cannot stop.

Using an Iterator

We could possibly take an iterator approach and write something like this:

class PrimeIt
{
public:    
    PrimeIt() : m_num ( 0 ) {}
    uint64_t next() {
        do {
            ++m_num;
        } while (!isPrime(m_num));
        return m_num;
    }
private:
    uint64_t m_num;
};

int main()
{
    PrimeIt it;
    for (int i = 0; i < 100; ++i) {
        std::cout << it.next() << " is prime\n";
    }
    return 0;
}

Except now we no longer have a nice simple to follow function. Where before the state (the latest prime number) was stored as a local variable, now it becomes a class member.

This is a simple example. Were we to resume from a nested function, all that state would have to be dragged out of our function and into the class.

This does solve our pause and resume problem though.

The Coroutine Generator

Coroutines allow us to write our prime number generator like in the Simple Case but use it like the Iterator case. It would look something like:

Generator<int64_t> getPrimes()
{
    int64_t num = 0;
    while (true) {
        while (!isPrime(num)) {
            ++num;
        }
        co_yield num;
        ++num;
    }
    co_return 0;
}


int main()
{
    Generator<int64_t> primes = getPrimes();
    for (int i = 0; i < 100; ++i) {
        std::cout << primes.next() << " is prime\n";
    }
    return 0;
}

You will notice that our main function is very similar to our Iterator case but getPrimes is very similar to an infinite version of nPrimes above. The only different is that getPrimes now has the line co_yield num. This line halts the coroutine and allows primes.next() to return the current coroutine value. This coroutine version will require a coroutine support library that might look something like:

//This code will probably be a general purpose library

#include <experimental/coroutine>
#include <iostream>

using namespace std::experimental::coroutines_v1;

template<typename T>
struct Promise;
template<typename T>
class Generator;

template<typename T>
using Handle = std::experimental::coroutines_v1::coroutine_handle<Promise<T>>;

template<typename T>
struct Promise
{
    Promise() : val { T() } {}
    std::experimental::coroutines_v1::suspend_always initial_suspend() { return {}; }
    std::experimental::coroutines_v1::suspend_always final_suspend() { return {}; }
    Generator<T> get_return_object();
    void unhandled_exception() { abort(); }
    void return_value(T val) {
        this->val = val;
    }
    std::experimental::coroutines_v1::suspend_always yield_value(T val) {
        this->val = val;
        return {};
    }
    
    T val;
};

template<typename T>
class Generator
{
public:
        explicit Generator(Handle<T> handle)
        : m_handle (handle) 
    {}
    ~Generator() {
        if (m_handle) {
            m_handle.destroy();
        }
    }
    T next();
    using promise_type = Promise<T>;
private:
    Handle<T> m_handle;    
};

template<typename T>
Generator<T> Promise<T>::get_return_object()
{
    return Generator<T> { Handle<T>::from_promise(*this) };
}


template<typename T>
T Generator<T>::next()
{
    if (m_handle) {
        m_handle.resume();
        return m_handle.promise().val;
    } else {
        return T();
    }
}

//End general purpose library

The Generator listed above is now templated and written once should be able to support a number of coroutine generators.

At this point I would like to mention the lines:

    std::experimental::coroutines_v1::suspend_always yield_value(T val) {
        this->val = val;
        return {};
    }

promise_type::yield_value is the function called by the compiler to return the value for the co_yield expression. The promise can then decide whether to suspend the coroutine or allow it to continue.

Now since C++ implemented stackless coroutines, a coroutine cannot store a large function callstack. Coroutines can however be chained. This and more will be discussed in future installments.