Compiler Explorer

Source code

#include <experimental/coroutine>
#include <cassert>
#include <fmt/core.h>

// forward decl some promise types
template<typename U> struct promise_type_impl_base;
template<typename U> struct promise_type_impl;

template<typename T>
class task {
public:
  using promise_type = promise_type_impl<T>;
  using handle_type = std::experimental::coroutine_handle<promise_type>;

task(handle_type handle) : handle_(handle) {  }
  task(task&) = delete; task(task&&) = delete; // disable copy/move
  bool await_ready() { return handle_.done(); }
  bool await_resume() {
    if (!handle_.done())
      handle_.resume();
    return !handle_.done();
  }
  template<typename PROMISE>
  void await_suspend(std::experimental::coroutine_handle<PROMISE> coroutine) {}
  ~task() { handle_.destroy(); }

private:
  handle_type handle_;
};

template<typename U>
struct promise_type_impl_base {
  auto get_return_object() { return task<U>::handle_type::from_promise(*static_cast<typename task<U>::promise_type*>(this));}

auto initial_suspend() noexcept { return std::experimental::suspend_always{}; }
  auto final_suspend() noexcept { return std::experimental::suspend_always{}; }
  void unhandled_exception() { std::terminate(); };

#if 0
  static void* operator new (std::size_t size, const std::nothrow_t& tag) noexcept {
	void * rv = malloc(size)/*nullptr*/; // set nullptr to hit get_return_object_on_allocation_failure 
	fmt::print("new {} at {}\n", size, rv);
	return rv;
  }
  static void operator delete (void* ptr) { fmt::print("delete {}\n", ptr); free(ptr); }
  static task<U> get_return_object_on_allocation_failure() { throw std::bad_alloc(); };
#endif
};

template<typename U>
struct promise_type_impl : public promise_type_impl_base<U> {    
  void return_value(U v) { _v = v; }
  U& result() { return _v; }
  // void yield_value(U v) { _v = v; }
  U _v;
};

template<>
struct promise_type_impl<void> : public promise_type_impl_base<void> {
  void return_void() {}
};

namespace stdx = std::experimental;

struct my_fake_coro_frame {
  using cb = void (*)(void *);
  cb resume, destroy;
  task<int>::promise_type p;
  int a;
  std::string s;
  char suspend_point_id;       // offset 48, suspension point index
  stdx::suspend_always initial_suspend_object; // offset 49, initial_suspend_object
  stdx::suspend_always suspend_point_1_object; // offset 50, suspend_point_1_object
  stdx::suspend_always suspend_point_2_object; // offset 51, suspend_point_1_object
  stdx::suspend_always final_suspend_object;   // offset 52, final_suspend_object
  char _53;                     // offset 53, padding
  char _54;                     // offset 54 ...
  char _55;                     // offset 55 ...
};
void my_fake_coro_resume(void *c);
void my_fake_coro_destroy(void *c);

task<int> my_fake_coro(int a) {
  // Allocate space on the heap for coroutine's frame -- todo: alignment?
  auto b = new char[sizeof(my_fake_coro_frame)];
  my_fake_coro_frame *f = (my_fake_coro_frame *)b;

// Set up the prefix internals
  f->resume = my_fake_coro_resume;
  f->destroy = my_fake_coro_destroy;

// Save the paramaters
  f->a = std::move(a);

// Set up the locals??
  // new (&f->s) std::string();
  // new (&f->p) task<int>::promise_type();

// Initalize suspension point id
  f->suspend_point_id = 0;

// Grab our return object
  auto ro = f->p.get_return_object();

f->initial_suspend_object = f->p.initial_suspend();
  if (not f->initial_suspend_object.await_ready()) {
    f->initial_suspend_object.await_suspend(
        stdx::coroutine_handle<task<int>::promise_type>::from_address(
            (void *)f));
    return ro;
  }

// Pump the state machine
  my_fake_coro_resume(b);

// Return in a task
  return ro;
}

void my_fake_coro_resume(void *c) {
  my_fake_coro_frame *f = (my_fake_coro_frame *)c;
  fmt::print("RESUME FAKE {} suspend point {}\n", (uintptr_t)c, (int)f->suspend_point_id);
  try {
    auto &s = f->s;

switch (f->suspend_point_id) {
    case 0:
      f->initial_suspend_object.await_resume();

s = "In";

f->suspend_point_id = 1;
      {
        f->suspend_point_1_object = stdx::suspend_always{};
        if (not f->suspend_point_1_object.await_ready()) {
          f->suspend_point_1_object.await_suspend(
              stdx::coroutine_handle<task<int>::promise_type>::from_address(c));
          return;
        }
      }
      [[fallthrough]];
    case 1:
      f->suspend_point_1_object.await_resume();

fmt::print(s);
      s = " coroutine!\n";

f->suspend_point_id = 2;
      {
        f->suspend_point_2_object = stdx::suspend_always{};
        if (not f->suspend_point_2_object.await_ready()) {
          f->suspend_point_2_object.await_suspend(
              stdx::coroutine_handle<task<int>::promise_type>::from_address(c));
          return;
        }
      }
      [[fallthrough]];
    case 2:
      f->suspend_point_2_object.await_resume();
      
      fmt::print(s);
      f->p.return_value(f->a);

f->suspend_point_id = 3;
      {
        f->final_suspend_object = stdx::suspend_always{};
        if (not f->final_suspend_object.await_ready()) {
          f->final_suspend_object.await_suspend(
              stdx::coroutine_handle<task<int>::promise_type>::from_address(c));
          return;
        }
      }
      [[fallthrough]];
    case 3:
      f->final_suspend_object.await_resume();
      f->resume = nullptr;
      break;
    }
  } catch (...) {
    f->p.unhandled_exception();
  }
}

void my_fake_coro_destroy(void *c) {
  my_fake_coro_frame *f = (my_fake_coro_frame *)c;
  delete f;

fmt::print("DESTROY FAKE {}\n", (uintptr_t)c);
}

task<int> test_coro(int a) {
  std::string s = "In";
  co_await std::experimental::suspend_always{};
  fmt::print(s);
  s = " coroutine!\n";
  co_await std::experimental::suspend_always{};
  fmt::print(s);
  co_return a;
}

int main(int argc, char **argv) {
    // Create and pump coroutine to completion
    
    auto fake = my_fake_coro(0);
    while(fake.await_resume());

auto test = test_coro(0);
    while(test.await_resume());
}