2013-09-18

Python-style generators in C++ (part 2)

Python-style generators in C++ (part 2)

In part 2 I will continue turn the code from part 1 into something more usable.

c++ implementation, phase 2: getting rid of globals

In Python, generator functions returns generator object, which encapsulates the private context and allows for controlling the execution. Nothing prevents the user to have more than one generator using the same function, running in parallel:
#!/usr/bin/env python

def fibonacci():
    last = 1
    current = 1
    while(True):
        yield current
        nxt = last + current
        last = current
        current = nxt


N = 10
print('Two fibonacci sequences generated in parallel:')
generator1 = fibonacci()
print('seq #1: %d' % generator1.next())
print('seq #1: %d' % generator1.next())

generator2 = fibonacci()
for i in range(N):
    print('seq #1: %d' % generator1.next())
    print('seq #2:       %d' % generator2.next())

Output:
seq #1: 1
seq #1: 2
seq #1: 3
seq #2:       1
seq #1: 5
seq #2:       2
seq #1: 8
seq #2:       3
seq #1: 13
seq #2:       5
seq #1: 21
seq #2:       8
seq #1: 34
seq #2:       13
seq #1: 55
seq #2:       21
seq #1: 89
seq #2:       34
seq #1: 144
seq #2:       55
seq #1: 233
seq #2:       89
In our C++ code, all global functions and variables need to be wrapped into a class. Unfortunately - that means that yield can no longer be global. It can be passed as a parameter to the generator function, but this will change the function signature. To make it work, the function call has to be wrapped; this is an perfect opportunity to solve another major problem: the generator function can now return or throw exception.
New C++11 exception tools can be used to capture any exception from the generator function and re-throw it in the main context.
// exception thrown where generator function exists
struct generator_finished : public std::exception
{
    virtual const char* what() const noexcept { return "generator finished"; }
};

class generator
{
public:

    typedef std::function<void(intptr_t)> yield_function_type;
    typedef std::function<void(yield_function_type)> generator_function_type;

    // former 'init()'
    generator(generator_function_type generator, std::size_t stack_size = DEFAULT_STACK_SIZE)
        : _generator(std::move(generator))
    {
        // allocate stack for new context
        _stack = new char[stack_size];

        // make a new context. The returned fcontext_t is created on the new stack, so there is no need to delete it
        _new_context = boost::context::make_fcontext(
                    _stack + stack_size, // new stack pointer. On x86/64 it hast be the TOP of the stack (hence the "+ STACK_SIZE")
                    stack_size,
                    &generator::static_generator_function); // will call generator wrapper
    // prevent copying
    generator(const generator&) = delete;
    // former 'cleanup()'
    ~generator()
    {
        delete _stack;
        _stack = nullptr;
        _new_context = nullptr;
    }

    intptr_t next()
    {
        // prevent calling when the generator function already finished
        if (_exception)
            std::rethrow_exception(_exception);

        // switch to function context. May set _exception
        intptr_t result = boost::context::jump_fcontext(&_main_context, _new_context, reinterpret_cast<intptr_t>(this));
        if (_exception)
            std::rethrow_exception(_exception);
        else
            return result;
    }

private:

    // former global variables
    boost::context::fcontext_t _main_context; // will hold the main execution context
    boost::context::fcontext_t* _new_context = nullptr; // will point to the new context
    static const int DEFAULT_STACK_SIZE= 64*1024; // completely arbitrary value
    char* _stack = nullptr;

    generator_function_type _generator; // generator function

    std::exception_ptr _exception = nullptr;// pointer to exception thrown by generator function


    // the actual generator function used to create context
    static void static_generator_function(intptr_t param)
    {
        generator* _this = reinterpret_cast<generator*>(param);
        _this->generator_wrapper();
    }

    void yield(intptr_t value)
    {
        boost::context::jump_fcontext(_new_context, &_main_context, value); // switch back to the main context
    }

    void generator_wrapper()
    {
        try
        {
            _generator([this](intptr_t value) // use lambda to bind this to yield
            {
                yield(value);
            });
            throw generator_finished();
        }
        catch(...)
        {
            // store the exception, is it can be thrown back in the main context
            _exception = std::current_exception();
            boost::context::jump_fcontext(_new_context, &_main_context, 0); // switch back to the main context
        }
    }
};

void fibonacci(const generator::yield_function_type& yield)
{
    intptr_t last = 1;
    intptr_t current = 1;
    for(;;)
    {
        yield(current);
        intptr_t nxt = last + current;
        last = current;
        current = nxt;
    }
}

int main(int , char** )
{
    const int N = 10;
    std::cout << "Two fibonacci sequences generated in parallel::" << std::endl;
    generator generator1(fibonacci);
    std::cout << "seq #1: " << generator1.next() << std::endl;
    std::cout << "seq #1: " << generator1.next() << std::endl;

    generator generator2(fibonacci);
    for(int i = 0; i < N; i++)
    {
        std::cout << "seq #1: " << generator1.next() << std::endl;
        std::cout << "seq #2:       " << generator2.next() << std::endl;
    }
}

Final touch: sprinkling the code with templates

The generator is almost ready, and now not only the fibonacci, but also main looks almost like their python counterparts.
The last remaining flaw to fix is the return type: current implementation can not return anything else than intptr_t. This can be solved by turning the generator into a template. The returns value can no longer be passed via jump_fcontext, but we can pass it by member variable, just like the exception.
#include <boost/context/all.hpp>
#include <boost/optional.hpp>
#include <iostream>
#include <iomanip>
#include <functional>
#include <exception>

// exception thrown where generator function exists
struct generator_finished : public std::exception
{
    virtual const char* what() const noexcept { return "generator finished"; }
};

template<typename ReturnType>
class generator
{
public:

    typedef std::function<void(const ReturnType&)> yield_function_type;
    typedef std::function<void(yield_function_type)> generator_function_type;

    // former 'init()'
    generator(generator_function_type generator, std::size_t stack_size = DEFAULT_STACK_SIZE)
        : _generator(std::move(generator))
    {
        // allocate stack for new context
        _stack = new char[stack_size];

        // make a new context. The returned fcontext_t is created on the new stack, so there is no need to delete it
        _new_context = boost::context::make_fcontext(
                    _stack + stack_size, // new stack pointer. On x86/64 it hast be the TOP of the stack (hence the "+ STACK_SIZE")
                    stack_size,
                    &generator::static_generator_function); // will call generator wrapper
    }

    // prevent copying
    generator(const generator&) = delete;

    // former 'cleanup()'
    ~generator()
    {
        delete _stack;
        _stack = nullptr;
        _new_context = nullptr;
    }

    ReturnType next()
    {
        // prevent calling when the generator function already finished
        if (_exception)
            std::rethrow_exception(_exception);

        // switch to function context. May set _exception
        boost::context::jump_fcontext(&_main_context, _new_context, reinterpret_cast<intptr_t>(this));
        if (_exception)
            std::rethrow_exception(_exception);
        else
            return *_return_value;
    }

private:

    // former global variables
    boost::context::fcontext_t _main_context; // will hold the main execution context
    boost::context::fcontext_t* _new_context = nullptr; // will point to the new context
    static const int DEFAULT_STACK_SIZE= 64*1024; // completely arbitrary value
    char* _stack = nullptr;

    generator_function_type _generator; // generator function

    std::exception_ptr _exception = nullptr;// pointer to exception thrown by generator function
    boost::optional<ReturnType> _return_value; // optional allows for using typed without defautl constructor


    // the actual generator function used to create context
    static void static_generator_function(intptr_t param)
    {
        generator* _this = reinterpret_cast<generator*>(param);
        _this->generator_wrapper();
    }

    void yield(const ReturnType& value)
    {
        _return_value = value;
        boost::context::jump_fcontext(_new_context, &_main_context, 0); // switch back to the main context
    }

    void generator_wrapper()
    {
        try
        {
            _generator([this](const ReturnType& value) // use lambda to bind this to yield
            {
                yield(value);
            });
            throw generator_finished();
        }
        catch(...)
        {
            // store the exception, is it can be thrown back in the main context
            _exception = std::current_exception();
            boost::context::jump_fcontext(_new_context, &_main_context, 0); // switch back to the main context
        }
    }
};

template<typename NumericType> // now we can choose in which flavour do we want our fibonacci numbers
void fibonacci(const typename generator<NumericType>::yield_function_type& yield)
{
    NumericType last = 1;
    NumericType current = 1;
    for(;;)
    {
        yield(current);
        NumericType nxt = last + current;
        last = current;
        current = nxt;
    }
}

int main(int , char** )
{
    const int N = 10;
    std::cout << "Two fibonacci sequences generated in parallel::" << std::endl;
    std::cout << std::setprecision(3) << std::fixed; // to make floating point number distinguishable
    generator<int> generator1(fibonacci<int>);
    std::cout << "seq #1: " << generator1.next() << std::endl;
    std::cout << "seq #1: " << generator1.next() << std::endl;

    generator<double> generator2(fibonacci<double>);
    for(int i = 0; i < N; i++)
    {
        std::cout << "seq #1: " << generator1.next() << std::endl;
        std::cout << "seq #2:       "  << generator2.next() << std::endl;
    }
}
Output:
Two fibonacci sequences generated in parallel::
seq #1: 1
seq #1: 2
seq #1: 3
seq #2:       1.000
seq #1: 5
seq #2:       2.000
seq #1: 8
seq #2:       3.000
seq #1: 13
seq #2:       5.000
seq #1: 21
seq #2:       8.000
seq #1: 34
seq #2:       13.000
seq #1: 55
seq #2:       21.000
seq #1: 89
seq #2:       34.000
seq #1: 144
seq #2:       55.000
seq #1: 233
seq #2:       89.000

What else?

Further work on the generator could include:
  • Passing arguments to generator function. This could be done with variadic templates and tuples
  • Implementing copy constructor. This would allow for taking snapshot of function execution state at any moment.
  • Making the code thread-aware. Currently generator::next() can be called from different threads, effectively making different parts of the function to run in different threads. Wicked!

No comments:

Post a Comment