simple/naive implementation of continuation/coroutine/generator on GNU/Linux


1 continuation

In some senses, ucontext.h is already a implementation of continuation, which is available on most modern OS(standarded by POSIX). In C99, we also have setjmp/longjmp to implement continuation. But setjmp/longjmp is weaker than ucontext.h.
Folowing code are just a demo, not suit for production usage. g++-4.7 is required.
#include 
#include  // for type erasure, or, we have to make continuation a template

#define CLOG(...) // we don't have/need CLOG

struct continuation {
  template <class Func>
  continuation(const Func & func, size_t stack_size=1<<18) {
    m_stack.resize(stack_size);
    m_func = func;
    int r = getcontext(&m_ctxs[1]);
    assert(r == 0);
    m_ctxs[1].uc_link = &m_ctxs[0];
    m_ctxs[1].uc_stack.ss_sp = &m_stack[0];
    m_ctxs[1].uc_stack.ss_size = m_stack.size();
    // TODO: x86_64 is safe to pass pointer with glibc, but for other platform ...
    makecontext(&m_ctxs[1], (void (*)(void))&s_context_func, 1, this);
  }

  continuation(const continuation &r) = delete;
  continuation & operator=(const continuation &r) = delete;

  int swap() {
    if (m_in_done) return -1;
    CLOG(DEBUG, "swap from", (m_is_in ? "in" : "out"));
    m_is_in = !m_is_in;
    swapcontext(&m_ctxs[!m_is_in], &m_ctxs[m_is_in]);
    CLOG(DEBUG, "after swap");
    if (m_in_done) return 0;
    return 1;
  }

  static void s_context_func(continuation *self) {
    return self->context_func();
  }

  void context_func() {
    m_func();
    m_in_done = 1;
  }

  ucontext_t m_ctxs[2];
  std::function<void(void)> m_func;
  std::vector<char> m_stack;
  bool m_is_in = false;
  int m_in_done = 0;
};

2 generator

template <class T>
struct generator : continuation {
  template <class Func>
  generator(const Func &func)
  : continuation(func) {

  }

  void yield(const T &value) {
    m_value = value;
    swap();
  }

  T yield() {
    if (swap() <= 0) throw std::runtime_error("end of iteration");
    return m_value;
  }

  T volatile m_value;
};

3 coroutine

struct coroutine : continuation {
  template <class Func>
  generator(const Func &func)
  : continuation(func) {

  }

  void yield(coroutine &other) {
    other.swap();
  }
};

4 usage and test

4.1 primes

A simple program to print primes(not all).
bool is_prime(int n) {
  if (n % 2 == 0) return n == 2;
  for (int p = 3; p * p <= n; p += 2) {
    if (n % p == 0) return false;
  }
  return true;
}

void primes(generator<int> &g) {
  g.yield(2);
  int i = 3;
  while (1) {
    if (is_prime(i)) g.yield(i);
    i += 2;
  }
}

int main() {
  generator<int> g([&]() {
      primes(g);
    });
  while (1) {
    int p = g.yield();
    std::cout << p << std::endl;
  }
}

4.2 producer and consumer

void producer(int n, int *x, continuation & cont) {
  for (int i = 0; i < n; ++i) {
    *x = i * i;
    cont.swap();
  }
}

void consumer(int volatile *x, continuation &cont) {
  while (cont.swap() > 0) {
    CLOG(INFO, "got", *x);
  }
}

No comments: