Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Making python generator via c++20 coroutines

Let's say I have this python code:

def double_inputs():
    while True:
        x = yield
        yield x * 2
gen = double_inputs()
next(gen)
print(gen.send(1))

It prints "2", just as expected. I can make a generator in c++20 like that:

#include <coroutine>

template <class T>
struct generator {
    struct promise_type;
    using coro_handle = std::coroutine_handle<promise_type>;

    struct promise_type {
        T current_value;
        auto get_return_object() { return generator{coro_handle::from_promise(*this)}; }
        auto initial_suspend() { return std::suspend_always{}; }
        auto final_suspend() { return std::suspend_always{}; }
        void unhandled_exception() { std::terminate(); }
        auto yield_value(T value) {
            current_value = value;
            return std::suspend_always{};
        }
    };

    bool next() { return coro ? (coro.resume(), !coro.done()) : false; }
    T value() { return coro.promise().current_value; }

    generator(generator const & rhs) = delete;
    generator(generator &&rhs)
        :coro(rhs.coro)
    {
        rhs.coro = nullptr;
    }
    ~generator() {
        if (coro)
            coro.destroy();
    }
private:
    generator(coro_handle h) : coro(h) {}
    coro_handle coro;
};

generator<char> hello(){
    //TODO:send string here via co_await, but HOW???
    std::string word = "hello world";
    for(auto &ch:word){
        co_yield ch;
    }
}

int main(int, char**) {
    for (auto i = hello(); i.next(); ) {
        std::cout << i.value() << ' ';
    }
}

This generator just produces a string letter by letter, but the string is hardcoded in it. In python, it is possible not only to yield something FROM the generator but to yield something TO it too. I believe it could be done via co_await in C++.

I need it to work like this:

generator<char> hello(){
    std::string word = co_await producer; // Wait string from producer somehow 
    for(auto &ch:word){
        co_yield ch;
    }
}

int main(int, char**) {
    auto gen = hello(); //make consumer
    producer("hello world"); //produce string
    for (; gen.next(); ) {
        std::cout << gen.value() << ' '; //consume string letter by letter
    }
}

How can I achieve that? How to make this "producer" using c++20 coroutines?

like image 654
tort_dla_psa Avatar asked Sep 26 '20 22:09

tort_dla_psa


1 Answers

You have essentially two problems to overcome if you want to do this.

The first is that C++ is a statically typed language. This means that the types of everything involved need to be known at compile time. This is why your generator type needs to be a template, so that the user can specify what type it shepherds from the coroutine to the caller.

So if you want to have this bi-directional interface, then something on your hello function must specify both the output type and the input type.

The simplest way to go about this is to just create an object and pass a non-const reference to that object to the generator. Each time it does a co_yield, the caller can modify the referenced object and then ask for a new value. The coroutine can read from the reference and see the given data.

However, if you insist on using the future type for the coroutine as both output and input, then you need to both solve the first problem (by making your generator template take OutputType and InputType) as well as this second problem.

See, your goal is to get a value to the coroutine. The problem is that the source of that value (the function calling your coroutine) has a future object. But the coroutine cannot access the future object. Nor can it access the promise object that the future references.

Or at least, it can't do so easily.

There are two ways to go about this, with different use cases. The first manipulates the coroutine machinery to backdoor a way into the promise. The second manipulates a property of co_yield to do basically the same thing.

Transform

The promise object for a coroutine is usually hidden and inaccessible from the coroutine. It is accessible to the future object, which the promise creates and which acts as an interface to the promised data. But it is also accessible during certain parts of the co_await machinery.

Specifically, when you perform a co_await on any expression in a coroutine, the machinery looks at your promise type to see if it has a function called await_transform. If so, it will call that promise object's await_transform on every expression you co_await on (at least, in a co_await that you directly write, not implicit awaits, such as the one created by co_yield).

As such, we need to do two things: create an overload of await_transform on the promise type, and create a type whose sole purpose is to allow us to call that await_transform function.

So that would look something like this:

struct generator_input {};

...

//Within the promise type:
auto await_transform(generator_input);

One quick note. The downside of using await_transform like this is that, by specifying even one overload of this function for our promise, we impact every co_await in any coroutine that uses this type. For a generator coroutine, that's not very important, since there's not much reason to co_await unless you're doing a hack like this. But if you were creating a more general mechanism that could distinctly await on arbitrary awaitables as part of its generation, you'd have a problem.

OK, so we have this await_transform function; what does this function need to do? It needs to return an awaitable object, since co_await is going to await on it. But the purpose of this awaitable object is to deliver a reference to the input type. Fortunately, the mechanism co_await uses to convert the awaitable into a value is provided by the awaitable's await_resume method. So ours can just return an InputType&:

//Within the `generator<OutputType, InputType>`:
    struct passthru_value
    {
        InputType &ret_;

        bool await_ready() {return true;}
        void await_suspend(coro_handle) {}
        InputType &await_resume() { return ret_; }
    };


//Within the promise type:
auto await_transform(generator_input)
{
    return passthru_value{input_value}; //Where `input_value` is the `InputType` object stored by the promise.
}

This gives the coroutine access to the value, by invoking co_await generator_input{};. Note that this returns a reference to the object.

The generator type can easily be modified to allow the ability to modify an InputType object stored in the promise. Simply add a pair of send functions for overwriting the input value:

void send(const InputType &input)
{
    coro.promise().input_value = input;
} 

void send(InputType &&input)
{
    coro.promise().input_value = std::move(input);
} 

This represents an asymmetric transport mechanism. The coroutine retrieves a value at a place and time of its own choosing. As such, it is under no real obligation to respond instantly to any changes. This is good in some respects, as it allows a coroutine to insulate itself from deleterious changes. If you're using a range-based for loop over a container, that container cannot be directly modified (in most ways) by the outside world or else your program will exhibit UB. So if the coroutine is fragile in that way, it can copy the data from the user and thus prevent the user from modifying it.

All in all, the needed code isn't that large. Here's a run-able example of your code with these modifications:

#include <coroutine>
#include <exception>
#include <string>
#include <iostream>

struct generator_input {};


template <typename OutputType, typename InputType>
struct generator {
    struct promise_type;
    using coro_handle = std::coroutine_handle<promise_type>;

    struct passthru_value
    {
        InputType &ret_;

        bool await_ready() {return true;}
        void await_suspend(coro_handle) {}
        InputType &await_resume() { return ret_; }
    };

    struct promise_type {
        OutputType current_value;
        InputType input_value;


        auto get_return_object() { return generator{coro_handle::from_promise(*this)}; }
        auto initial_suspend() { return std::suspend_always{}; }
        auto final_suspend() { return std::suspend_always{}; }
        void unhandled_exception() { std::terminate(); }
        auto yield_value(OutputType value) {
            current_value = value;
            return std::suspend_always{};
        }

        void return_void() {}

        auto await_transform(generator_input)
        {
            return passthru_value{input_value};
        }
    };

    bool next() { return coro ? (coro.resume(), !coro.done()) : false; }
    OutputType value() { return coro.promise().current_value; }

    void send(const InputType &input)
    {
        coro.promise().input_value = input;
    } 

    void send(InputType &&input)
    {
        coro.promise().input_value = std::move(input);
    } 

    generator(generator const & rhs) = delete;
    generator(generator &&rhs)
        :coro(rhs.coro)
    {
        rhs.coro = nullptr;
    }
    ~generator() {
        if (coro)
            coro.destroy();
    }
private:
    generator(coro_handle h) : coro(h) {}
    coro_handle coro;
};

generator<char, std::string> hello(){
    auto word = co_await generator_input{};

    for(auto &ch: word){
        co_yield ch;
    }
}

int main(int, char**)
{
    auto test = hello();
    test.send("hello world");

    while(test.next())
    {
        std::cout << test.value() << ' ';
    }
}

Be more yielding

An alternative to using an explicit co_await is to exploit a property of co_yield. Namely, co_yield is an expression and therefore it has a value. Specifically, it is (mostly) equivalent to co_await p.yield_value(e), where p is the promise object (ohh!) and e is what we're yielding.

Fortunately, we already have a yield_value function; it returns std::suspend_always. But it could also return an object that always suspends, but also which co_await can unpack into an InputType&:

struct yield_thru
{
    InputType &ret_;

    bool await_ready() {return false;}
    void await_suspend(coro_handle) {}
    InputType &await_resume() { return ret_; }
};

...

//in the promise
auto yield_value(OutputType value) {
    current_value = value;
    return yield_thru{input_value};
}

This is a symmetric transport mechanism; for every value you yield, you receive a value (which may be the same one as before). Unlike the explicit co_await method, you can't receive a value before you start to generate them. This could be useful for certain interfaces.

And of course, you could combine them as you see fit.

like image 179
Nicol Bolas Avatar answered Oct 23 '22 09:10

Nicol Bolas