Glue for using C++ coroutines with Panda3D's task system

async_coro.hpp

#include <coroutine>

class AsyncCoroutineTask final : public AsyncTask {
public:
  AsyncCoroutineTask() = default;
  ~AsyncCoroutineTask() = default;
  ALLOC_DELETED_CHAIN(AsyncCoroutineTask);

  virtual DoneStatus do_task() {
    _done_status = DS_cont;
    _handle.resume();
    if (_handle.done()) {
      _handle.destroy();
      return DS_done;
    }
    return _done_status;
  }

  void suspend() {
    _done_status = DS_await;
  }

  DoneStatus _done_status;
  std::coroutine_handle<> _handle;
};

namespace std {
  template<class... Args>
  struct coroutine_traits<PT(AsyncTask), Args...> {
    struct promise_type {
      promise_type() : _task(new AsyncCoroutineTask) {}

      std::suspend_always initial_suspend() { return {}; }
      std::suspend_always final_suspend() { return {}; }

      AsyncTask *get_return_object() {
        _task->_handle = std::coroutine_handle<decltype(*this)>::from_promise(*this);
        return _task;
      }

      void return_void() {
        _task->set_result(nullptr);
      }

      template<class T>
      void return_value(T value) {
        _task->set_result(value);
      }

      std::suspend_always yield_value(AsyncTask::DoneStatus status) {
        _task->_done_status = status;
        return {};
      }

      auto await_transform(PT(AsyncFuture) future) {
        struct awaiter {
          awaiter(PT(AsyncFuture) fut) : _future(std::move(fut)) {}

          bool await_ready() {
            return _future->done();
          };

          void await_suspend(std::coroutine_handle<promise_type> handle) {
            AsyncCoroutineTask *task = handle.promise()._task;
            task->suspend();
            _future->add_waiting_task(task);
          }

          EventParameter await_resume() {
            TypedObject *ptr = _future->get_result();
            if (ptr == nullptr) {
              return {};
            }
            else if (ptr->is_of_type(TypedWritableReferenceCount::get_class_type())) {
              return (TypedWritableReferenceCount *)ptr;
            }
            else if (ptr->is_of_type(TypedReferenceCount::get_class_type())) {
              return (TypedReferenceCount *)ptr;
            }
            return {};
          }

          PT(AsyncFuture) _future;
        };
        if (future->is_task() && !future->done()) {
          AsyncTask *task = (AsyncTask *)future.p();
          if (!task->is_alive()) {
            _task->get_manager()->add(task);
          }
        }
        return awaiter { std::move(future) };
      }
      auto await_transform(std::suspend_never arg) { return arg; }

      PT(AsyncCoroutineTask) _task;
    };
  };
};