Yielding Generators

Created 15th July, 2017 05:28 (UTC), last edited 30th July, 2017 03:59 (UTC)

We've seen how the promise_type together with the coroutine return type handles the interactions between the caller and the coroutine itself.

Our target is to be able to do something pretty simple:

generator count() {
    std::cout << "Going to yield 1" << std::endl;
    co_yield 1;
    std::cout << "Going to yield 2" << std::endl;
    co_yield 2;
    std::cout << "Going to yield 3" << std::endl;
    co_yield 3;
    std::cout << "count() is done" << std::endl;
}

co_yield is actually very similar to our co_return example previously. There is really just one additional component, and one small change:

  1. Because our generator doesn't have any sort of return value there is an implicit return that produces nothing. This means that instead of return_value we are going to need to fill in a return_void method.
  2. For each value that is yielded the compiler will arrange to make a call to yield_value.

These two things together mean that we need to suspend at the yields (so the user of the generator gets a chance to read the value) and we know we're finished yielding values when the return_void promise method gets called.

Let's get started

Let's sketch out something that should look fairly familiar:

struct generator {
    struct promise_type {
    };
    using handle_type =
        std::experimental::coroutine_handle<promise_type>;
private:
    generator(handle_type h)
    : coro(h) {
        std::cout << "generator" << std::endl;
    }
    handle_type coro;
};

This should look fairly similar to our lazy and sync examples.

We can track the lifetime of the promise to help us see what's going on and we also suspend the coroutine both at its start and end, all by adding these members to our promise_type:

promise_type() {
    std::cout << "Promise created" << std::endl;
}
~promise_type() {
    std::cout << "Promise died" << std::endl;
}
auto initial_suspend() {
    std::cout << "Started the coroutine, let's pause!" << std::endl;
    return std::experimental::suspend_always{};
}
auto final_suspend() {
    std::cout << "Finished the coroutine" << std::endl;
    return std::experimental::suspend_always{};
}

We still need to return the generator object from our promise_type, and we also have to do something if an exception happens.

auto get_return_object() {
    std::cout << "Send back a generator" << std::endl;
    return generator{handle_type::from_promise(*this)};
}
void unhandled_exception() {
    std::cout << "Exception...." << std::endl;
    std::exit(1);
}

For our co_return based lazy and sync classes we also had to provide a mechanism for the co_return value to find its way to the caller. We actually still have to do this, but our coroutine example doesn't have a return statement, so now we have to implement the return_void member:

auto return_void() {
    std::cout << "return_void" << std::endl;
    return std::experimental::suspend_never{};
}

We're not going to suspend here because we want to just carry on and suspend at the final suspension point (where we will stop so the generator can clean up for us).

Yielding values

The last thing we need is to handle the yield itself. When the coroutines gets to a co_yield the compiler will call yield_value on our promise_type so we can see the value and decide what to do.

auto yield_value(int value) {
    std::cout << "Yielded " << value << std::endl;
    current_value = value;
    return std::experimental::suspend_always{};
}

It's also important that we suspend here. If we don't then our coroutine will carry on and run another co_yield. This will cause yield_value to be called again and we'll lose the value that was just given to us—it'll get overwritten by the new one.

As we saw before, we now have the coroutine performing its half of the control flow. The rest of the control flow that happens needs to be in the generator code because that is what our caller will interact with.

There is a dance between the generator and the promise_type to control the resumption of the coroutine and generation of values in exactly the right way. We need to make sure that we suspend and resume so that we get all of the values and don't skip any (especially easy at the start and end).

We need some sort of API that users of generator are going to use to fetch values. The obvious one is some sort of iterator based interface, but let's start with something a little bit simpler that involves less code.

int main() {
    auto g = count();
    while ( g.move_next() )
        std::cout << ">> " << g.current_value() << std::endl;
    return 0;
}

The implication here is that move_next will resume the coroutine until a co_yield or a return happens. If a co_yield happens we're going to return true, if the coroutine returns then we're going to return false to signal the end of the yielded sequence. Inside the loop we expect the current_value method to return the last value yielded.

The latter is very straight forward:

int current_value() {
    return coro.promise().current_value;
}

This is exactly as we had with the earlier sync and lazy examples. We store the last yielded value in the promise and can easily retrieve it through the coroutine handle we store in the generator.

For the sequence for fetching values we want this:

  1. Enter the coroutine, but suspend before doing anything.
  2. When the move_next is called resume the coroutine so that it can either return or yield a value.
  3. If the corotune is now done we didn't get a value so move_next returns false.
  4. If we're not done then return true so that move_next can be called again (going back to step 2).

Let's do that and also track what's going on:

bool move_next() {
    std::cout << "Moving to next" << std::endl;;
    coro.resume();
    std::cout << "Are we done? ";
    auto still_going = not coro.done();
    std::cout << (still_going ? "There is another" : "We're done") << std::endl;
    return still_going;
}

We now have enough to run:

Promise created
Send back a generator
generator
Started the coroutine, let's pause!
Moving to next
Going to yield 1
Yielded 1
Are we done? There is another
>> 1
Moving to next
Going to yield 2
Yielded 2
Are we done? There is another
>> 2
Moving to next
Going to yield 3
Yielded 3
Are we done? There is another
>> 3
Moving to next
count() is done
return_void
Finished the coroutine, brakes on
Are we done? We're done

Examining this we can see the dance between the two sides, the promise_type and the generator, play out in the right way.

  1. The promise creates and returns the generator.
  2. The coroutine is entered, but suspended.
  3. The generator's move_next is called and it advances the coroutine.
  4. The coroutine yields its first value
  5. move_next finds that the coroutine has yielded so returns true
  6. The value is printed
  7. Steps 3 to 7 are repeated another couple of times

The last time though is a bit different. The final move_next results in a call to return_void rather than yield_value and the coroutine ends up suspended at the final suspension point. coro.done() now returns true so we know that the coroutine has yielded everything it can (or wants to) and move_next returns false and we're all done.

To see the importance of getting this dance right let's change the initial_suspend to keep going:

auto initial_suspend() {
    std::cout << "Started the coroutine, keep going!" << std::endl;
    return std::experimental::suspend_never{};
}

Now we'll see this sequence:

Promise created
Send back a generator
generator
Started the coroutine, keep going!
Going to yield 1
Yielded 1
Moving to next
Going to yield 2
Yielded 2
Are we done? There is another
>> 2
Moving to next
Going to yield 3
Yielded 3
Are we done? There is another
>> 3
Moving to next
count() is done
return_void
Finished the coroutine, brakes on
Are we done? We're done

And because we didn't stop at the beginning we've resumed before the first value could be read. We can fix this by changing the way the generator is used:

int main() {
    std::cout << "coro3" << std::endl;
    auto g = count();
    do {
        std::cout << ">> " << g.current_value() << std::endl;
    } while ( g.move_next() );
    return 0;
}

Now we see all of our values:

Promise created
Send back a generator
generator
Started the coroutine, let's pause!
Going to yield 1
Yielded 1
>> 1
Moving to next
Going to yield 2
Yielded 2
Are we done? There is another
>> 2
Moving to next
Going to yield 3
Yielded 3
Are we done? There is another
>> 3
Moving to next
count() is done
return_void
Finished the coroutine, brakes on
Are we done? We're done

But of course we have another problem. The code that uses the generator now assumes that the generator must yield at least one value. Probably our first API is the best for a general use generator so let's put that back.

Fixing the leak

We still have a more subtle problem though. Notice we never see the “Promise died” message. This is because we're not managing the life time of the coroutine itself yet. We have the same issue that we saw with sync and lazy. We need to make the generator movable, but not copyable and then destroy the coroutine at the correct point in time.

~generator() {
    std::cout << "~generator "
        << (not coro ? "(empty)" : "(contains a coroutine)")
        << std::endl;
    if ( coro ) coro.destroy();
}
generator(const generator &) = delete;
generator(generator &&g)
: coro(g.coro) {
    g.coro = nullptr;
};

Now we can see the coroutine itself is destroyed when our generator goes out of scope:

Promise created
Send back a generator
generator
Started the coroutine, let's pause!
~generator (empty)
Moving to next
Going to yield 1
Yielded 1
Are we done? There is another
>> 1
Moving to next
Going to yield 2
Yielded 2
Are we done? There is another
>> 2
Moving to next
Going to yield 3
Yielded 3
Are we done? There is another
>> 3
Moving to next
count() is done
return_void
Finished the coroutine, brakes on
Are we done? We're done
~generator (contains a coroutine)
Promise died

Next we're going to convert this to an iterator, a straightforward enough procedure, but with a few things worth pointing out.


Categories: