Simple Emscripten coroutine scheduler

_Emscripten Coroutines.md

• Replace void SceneLoop() with TASK(void) SceneLoop()
• Replace AnotherSceneLoop() with AWAIT AnotherSceneLoop()
• Replace flip_screen() with AWAIT flip_screen()
• Replace return with CO_RETURN
• At the end of flip_screen() impl, AWAIT next_frame()

coro.cpp

#ifdef USE_COROUTINES
#include "coro.h"
#include <functional>
#include <map>
#include <queue>
#include <set>
#include <stack>

#undef main

#ifdef __EMSCRIPTEN__
#include <emscripten.h>
#endif  // __EMSCRIPTEN__

namespace coro {

struct executor {
	// the task which should wake on the next frame
	std::stack<coroutine_handle<>> awake_next_frame;

	// stack of tasks which are waiting to be launched
	std::stack<std::function<task<void>()>> launching;

	// map from awaitee to everything it should wake when done
	std::map<coroutine_handle<>, std::set<coroutine_handle<>>> wake_on_done;

	bool frame() {
		//printf(" frame: %lu\n", awake_next_frame.size());

		std::queue<coroutine_handle<>> awake;

		while (!launching.empty()) {
			launching.top()();
			launching.pop();
		}

		// dump awake_next_frame into awake
		if (!awake_next_frame.empty()) {
			awake.push(awake_next_frame.top());
			awake_next_frame.pop();
		}

		// while anything is awake, attempt to resume it
		while (!awake.empty()) {
			coroutine_handle<> current = awake.front();
			//printf("  resume %p\n", current.address());
			awake.pop();

			current.resume();
			// if resuming it finished it, awake its dependents
			if (current.done()) {
				//printf("  done\n");
				auto it = wake_on_done.find(current);
				if (it != wake_on_done.end()) {
					for (coroutine_handle<> h : it->second) {
						//printf("    wakes %p\n", h.address());
						awake.push(h);
					}
					wake_on_done.erase(it);
				}
				current.destroy();  // it's done, so destroy it
			}
#ifdef _DEBUG
			else {
				// if it isn't put to sleep, it's still awake
				bool asleep = false;
				for (const auto& pair : wake_on_done) {
					if (pair.second.find(current) != pair.second.end()) {
						asleep = true;
						break;
					}
				}
				if (!asleep) {
					printf("a task yielded without going to sleep\n");
					awake.push(current);
				}
			}
#endif
		}

		return !awake_next_frame.empty();
	}

	bool schedule(coroutine_handle<> awaiter, coroutine_handle<> awaitee) {
		//printf("schedule(): %p awaits on %p\n", awaiter.address(), awaitee.address());
		wake_on_done[awaitee].insert(awaiter);
		return true;  // return control to resume()r
	}

} g_executor;

struct flip_awaiter {
	bool await_ready() {
		//printf("  flip_awaiter::await_ready\n");
		return false;
	}

	void await_suspend(coroutine_handle<> h) {
		//printf("  flip_awaiter::await_suspend(%p)\n", h.address());
		g_executor.awake_next_frame.push(h);
	}

	void await_resume() {
		//printf("  flip_awaiter::await_resume\n");
	}
};

task<void> next_frame() {
	//printf("next_frame\n");
	co_await flip_awaiter{};
}

void launch(std::function<task<void>()> entry_point) {
	g_executor.launching.push(entry_point);
}

bool __schedule(coroutine_handle<> awaiter, coroutine_handle<> awaitee) {
	g_executor.schedule(awaiter, awaitee);
	return true;
}

}  // namespace coro

#ifdef __EMSCRIPTEN__

static void em_main_loop() {
	if (!coro::g_executor.frame()) {
		printf("calling emscripten_cancel_main_loop\n");
		emscripten_cancel_main_loop();
	}
}

int main(int argc, char** argv) {
	coro::launch([=]() -> coro::task<void> {
		AWAIT coro::main(argc, argv);
	});
	emscripten_set_main_loop(em_main_loop, 0, 1);
	return 0;
}

#else  // __EMSCRIPTEN__

int main(int argc, char** argv) {
	int coro_retval = 0;
	coro::launch([&]() -> coro::task<void> {
		coro_retval = AWAIT coro::coro_main(argc, argv);
	});
	while (coro::g_executor.frame()) {}
	return coro_retval;
}

#endif  // __EMSCRIPTEN__

#endif // USE_COROUTINES

coro.h

#ifndef HAM_CORO_H
#define HAM_CORO_H

#include <functional>

#ifdef USE_COROUTINES
#include "coro_active.inc.h"
#else  // USE_COROUTINES

namespace coro {

inline void next_frame() {}
inline void launch(std::function<void()> entry_point) { entry_point(); }
inline int main() { return 0; }

};  // namespace coro

#define TASK(TY) TY
#define AWAIT
#define CO_RETURN return

#endif  // USE_COROUTINES
#endif  // HAM_CORO_H

coro_active.inc.h

// The meat of the coroutine header. In a separate file so that changes to it
// don't cause recompiles on platforms where it isn't in use.

#ifndef __cpp_coroutines
#error "Asked for coroutines, but there's not compiler support."
#endif

#include <stdio.h>
#include <experimental/coroutine>

namespace coro {

using namespace std::experimental::coroutines_v1;

bool __schedule(coroutine_handle<> awaiter, coroutine_handle<> awaitee);

template<typename T>
struct shared_state {
	bool filled = false;
	T value;
};

template<>
struct shared_state<void> {
	bool filled = false;
};

template<typename Result>
class task_base {
public:
	coroutine_handle<> handle;
	std::shared_ptr<shared_state<Result>> state;

	void await_resume_checks() {
		if (!state) {
			printf("task_base::await_resume_checks: state ptr is null\n");
			//exit(1);
		}
		if (!state->filled) {
			printf("task_base::await_resume_checks: state is not filled\n");
			//exit(1);
		}
	}

public:
	task_base(
		coroutine_handle<> handle,
		std::shared_ptr<shared_state<Result>> state
	) : handle(handle), state(state) {}

	task_base(task_base&&) = default;
	task_base(const task_base&) = delete;
	task_base& operator=(task_base&&) = default;
	task_base& operator=(const task_base &) = delete;

	bool await_ready() {
		//printf("  task(%p)::await_ready()\n", handle.address());
		// return false to suspend or true to continue
		return state->filled;
	}

	bool await_suspend(coroutine_handle<> h) {
		//printf("  task(%p)::await_suspend(%p)\n", handle.address(), h.address());
		//return g_executor.schedule(h, handle);
		return __schedule(h, handle);
	}
};

template<typename Result>
struct task : public task_base<Result> {
	task(coroutine_handle<> handle, std::shared_ptr<shared_state<Result>> state)
		: task_base<Result>(handle, state) {}

	Result await_resume() {
		//printf("  task<T>(%p)::await_resume()\n", this);
		this->await_resume_checks();

		Result v = std::move(this->state->value);
		this->state.reset();
		return v;
	}
};

template<>
struct task<void> : public task_base<void> {
	task(coroutine_handle<> handle, std::shared_ptr<shared_state<void>> state)
		: task_base<void>(handle, state) {}

	void await_resume() {
		//printf("  task<void>(%p)::await_resume\n", this);
		await_resume_checks();
	}
};

template<typename Result>
class promise_base {
protected:
	std::shared_ptr<shared_state<Result>> state;
public:
	promise_base() : state(std::make_shared<shared_state<Result>>()) {}

	suspend_never initial_suspend() {
		//printf("  promise_base::initial_suspend()\n");
		return {};
	}

	suspend_always final_suspend() {
		//printf("  promise_base::final_suspend()\n");
		return {};
	}

	void unhandled_exception() {
		//printf("  promise_base::unhandled_exception()\n");
	}

	task<Result> get_return_object() {
		//printf("  promise<T>::get_return_object\n");
		return {
			coroutine_handle<promise_base<Result>>::from_promise(*this),
			this->state
		};
	}
};

template<typename Result>
struct promise : public promise_base<Result> {
	void return_value(Result value) {
		//printf("  promise<T>::return_value\n");
		this->state->filled = true;
		this->state->value = value;
	}
};

template<>
struct promise<void> : public promise_base<void> {
	void return_void() {
		//printf("  promise<void>::return_void\n");
		this->state->filled = true;
	}
};

task<void> next_frame();
void launch(std::function<task<void>()> entry_point);
task<int> main(int argc, char** argv);

}  // namespace coro

template<typename Result, typename ...Arg>
struct std::experimental::coroutine_traits<coro::task<Result>, Arg...> {
	typedef coro::promise<Result> promise_type;
};

#define TASK(TY) __attribute__((warn_unused_result)) ::coro::task<TY>
#define AWAIT co_await
#define CO_RETURN co_return

#define main coro::main