[feature] add await for WaitHandle

include/ilias/task/detail/view.hpp

@@ -56,7 +56,8 @@ class TaskView<detail::TaskViewNull> {
     auto resume() const noexcept { return mHandle.resume(); }
 
     /**
-     * @brief Schedule the task in the executor
+     * @brief Schedule the task in the executor (thread safe)
+     * @note Do not call this function if the task is already scheduled
      * 
      * @return auto 
      */
@@ -209,6 +210,17 @@ class TaskView final : public TaskView<> {
     operator handle_type() const {
         return handle();
     }
+
+    /**
+     * @brief Cast from TaskView<> (note it is dangerous)
+     * 
+     * @param view 
+     * @return TaskView<T> 
+     */
+    static auto cast(TaskView<detail::TaskViewNull> view) {
+        auto handle = std::coroutine_handle<>(view);
+        return TaskView<T>(handle_type::from_address(handle.address()));
+    }
 };
 
 ILIAS_NS_END

include/ilias/task/executor.hpp

@@ -23,10 +23,10 @@ ILIAS_NS_BEGIN
 class Executor {
 public:
     /**
-     * @brief Post a callable to the executor
+     * @brief Post a callable to the executor (thread safe)
      * 
-     * @param fn 
-     * @param args 
+     * @param fn The function to post (can not be null)
+     * @param args The arguments of the function
      */
     virtual auto post(void (*fn)(void *), void *args) -> void = 0;
 

include/ilias/task/spawn.hpp

@@ -1,9 +1,22 @@
+/**
+ * @file spawn.hpp
+ * @author BusyStudent (fyw90mc@gmail.com)
+ * @brief Task spawn, wait, and cancel.
+ * @version 0.1
+ * @date 2024-09-30
+ * 
+ * @copyright Copyright (c) 2024
+ * 
+ */
 #pragma once
 
 #include <ilias/task/detail/view.hpp>
 #include <ilias/task/task.hpp>
 #include <concepts>
 
+#define ilias_go    ILIAS_NS::detail::SpawnTags{ } << 
+#define ilias_spawn ILIAS_NS::detail::SpawnTags{ } <<
+
 ILIAS_NS_BEGIN
 
 namespace detail {
@@ -113,6 +126,8 @@ struct SpawnDataRef {
     auto operator ->() const noexcept { return mPtr; }
 
     auto operator <=>(const SpawnDataRef &other) const noexcept = default;
+
+    operator bool() const noexcept { return mPtr != nullptr; }
 private:
     auto ref() noexcept -> void {
         if (mPtr) {
@@ -128,6 +143,33 @@ struct SpawnDataRef {
     SpawnData *mPtr = nullptr;
 };
 
+/**
+ * @brief Awaiter for Wait Handle
+ * 
+ * @tparam T 
+ */
+template <typename T>
+class WaitHandleAwaiter {
+public:
+    WaitHandleAwaiter(TaskView<T> task) : mTask(task) { }
+
+    auto await_ready() const -> bool { return mTask.done(); }
+
+    auto await_suspend(TaskView<> caller) -> void {
+        mTask.setAwaitingCoroutine(caller); //< When the task is done, resume the caller
+        mReg = caller.cancellationToken().register_( //< Let the caller's cancel request cancel the current task
+            &CancelTheTokenHelper, &mTask.cancellationToken()
+        );
+    }
+
+    auto await_resume() const -> Result<T> {
+        return mTask.value();
+    }
+private:
+    CancellationToken::Registration mReg; //< The reg of we wait for cancel
+    TaskView<T> mTask;
+};
+
 /**
  * @brief Check if callable and args can be used to create a task.
  * 
@@ -140,6 +182,8 @@ concept TaskGenerator = requires(Callable &&callable, Args &&...args) {
     std::invoke(callable, args...)._leak();
 };
 
+struct SpawnTags { };
+
 } // namespace detail
 
 /**
@@ -148,6 +192,9 @@ concept TaskGenerator = requires(Callable &&callable, Args &&...args) {
  */
 class CancelHandle {
 public:
+    explicit CancelHandle(const detail::SpawnDataRef &data) : 
+        mData(data) { }
+
     CancelHandle() = default;
     CancelHandle(std::nullptr_t) { }
     CancelHandle(const CancelHandle &) = default;
@@ -157,6 +204,14 @@ class CancelHandle {
     auto cancel() const -> void { return mData->mTask.cancel(); }
     auto operator =(const CancelHandle &) -> CancelHandle & = default;
     auto operator <=>(const CancelHandle &) const noexcept = default;
+
+    /**
+     * @brief Check the CancelHandle is valid.
+     * 
+     * @return true 
+     * @return false 
+     */
+    operator bool() const { return bool(mData); }
 private:
     detail::SpawnDataRef mData;
 template <typename T>
@@ -182,24 +237,51 @@ class WaitHandle {
     auto cancel() const -> void { return mData->mTask.cancel(); }
 
     /**
-     * @brief Wait for the task to complete. and return the result.
+     * @brief Blocking Wait for the task to complete. and return the result.
      * 
      * @return Result<T> 
      */
     auto wait() -> Result<T> {
+        ILIAS_ASSERT_MSG(mData, "Can not wait for an invalid handle");
         if (!done()) {
             // Wait until done
             CancellationToken token;
             mData->mTask.registerCallback(detail::CancelTheTokenHelper, &token);
             mData->mTask.executor()->run(token);
         }
-        auto value = static_cast<TaskView<T> &>(mData->mTask).value();
+        auto value = TaskView<T>::cast(mData->mTask).value();
         mData = nullptr;
         return value;
     }
 
     auto operator =(const WaitHandle &) = delete;
     auto operator =(WaitHandle &&) -> WaitHandle & = default;
+    auto operator <=>(const WaitHandle &) const noexcept = default;
+
+    /**
+     * @brief co-await the handle to complete. and return the result.
+     * 
+     * @return co_await 
+     */
+    auto operator co_await() && {
+        ILIAS_ASSERT_MSG(mData, "Can not await an invalid handle");
+        return detail::WaitHandleAwaiter<T>{TaskView<T>::cast(mData->mTask)};
+    }
+
+    /**
+     * @brief Check if the handle is valid.
+     * 
+     * @return true 
+     * @return false 
+     */
+    operator bool() const { return bool(mData); }
+
+    /**
+     * @brief implicitly convert to CancelHandle
+     * 
+     * @return CancelHandle 
+     */
+    operator CancelHandle() const { return CancelHandle(mData); }
 private:
     detail::SpawnDataRef mData;
 };
@@ -215,7 +297,7 @@ template <typename T>
 inline auto spawn(Task<T> &&task) -> WaitHandle<T> {
     auto ref = detail::SpawnDataRef(new detail::SpawnData);
     ref->mTask = task._leak();
-    ref->mTask.executor()->schedule(ref->mTask); //< Start it on the event loop
+    ref->mTask.schedule(); //< Start it on the event loop
     return WaitHandle<T>(ref);
 }
 
@@ -247,9 +329,21 @@ inline auto spawn(Callable callable, Args &&...args) {
                 std::forward<Args>(args)...
             )
         );
-        ref->mTask.executor()->schedule(ref->mTask); //< Start it on the event loop
+        ref->mTask.schedule(); //< Start it on the event loop
         return WaitHandle<typename std::invoke_result_t<Callable, Args...>::value_type>(ref);
     }
 }
 
+/**
+ * @brief Helper for ilias_go macro and ilias_spawn macro
+ * 
+ * @tparam Args 
+ * @param args 
+ * @return auto 
+ */
+template <typename ...Args>
+inline auto operator <<(detail::SpawnTags, Args &&...args) {
+    return spawn(std::forward<Args>(args)...);
+}
+
 ILIAS_NS_END

include/ilias/task/task.hpp

@@ -179,25 +179,24 @@ class Task {
     explicit operator bool() const noexcept {
         return mHandle;
     }
+
+    /**
+     * @brief Get the awaiter of the task
+     * 
+     * @tparam T 
+     * @param task 
+     * @return detail::TaskAwaiter<T> 
+     */
+    auto operator co_await() && -> detail::TaskAwaiter<T> {
+        return detail::TaskAwaiter<T>(mHandle);
+    }
 private:
     Task(handle_type handle) : mHandle(handle) { }
 
     handle_type mHandle;
 friend class detail::TaskPromise<T>;
 };
 
-/**
- * @brief Get the awaiter of the task
- * 
- * @tparam T 
- * @param task 
- * @return detail::TaskAwaiter<T> 
- */
-template <typename T>
-inline auto operator co_await(Task<T> &&task) -> detail::TaskAwaiter<T> {
-    return detail::TaskAwaiter(task._view());
-}
-
 /**
  * @brief Sleep the current task for a period of time
  * 

include/ilias/zlib.hpp

@@ -14,6 +14,7 @@
 #if __has_include(<zlib.h>) && !defined(ILIAS_NO_ZLIB)
 
 #include <ilias/detail/expected.hpp>
+#include <ilias/task/task.hpp>
 #include <ilias/io/traits.hpp>
 #include <ilias/buffer.hpp>
 #include <ilias/error.hpp>

include/xmake.lua

@@ -2,5 +2,5 @@ set_languages("c++latest")
 
 target("ilias")
     set_kind("headeronly")
-    add_headerfiles("ilias/*.hpp")
+    add_headerfiles("ilias/**.hpp")
 target_end()

tests/unit/task/spawn.cpp

@@ -3,9 +3,9 @@
 #include <gtest/gtest.h>
 
 using namespace ILIAS_NAMESPACE;
+using namespace std::literals;
 
 TEST(TaskSpawn, SpawAndWait) {
-    MiniExecutor exec;
     auto callable = []() -> Task<int> {
         co_return 42;
     };
@@ -20,21 +20,40 @@ TEST(TaskSpawn, SpawAndWait) {
 }
 
 TEST(TaskSpawn, Detach) {
-    MiniExecutor exec;
     int value = 0;
     spawn([&]() -> Task<> {
         value = 1;
         co_return {};
     });
     spawn([&]() -> Task<> {
-        co_await exec.sleep(10);
+        co_await sleep(10ms);
         co_return {};
     }).wait();
 
     ASSERT_EQ(value, 1);
 }
 
+TEST(TaskSpawn, Await) {
+    auto handle = spawn([]() -> Task<int> {
+        co_return 42;
+    });
+    auto value = [&]() -> Task<int> {
+        co_return co_await std::move(handle);
+    }().wait().value();
+
+    ASSERT_EQ(value, 42);
+}
+
+TEST(TaskSpawn, Macro) {
+    auto fn = []() -> Task<int> {
+        co_return 42;
+    };
+    auto answer = ilias_go fn();
+    ASSERT_EQ(answer.wait().value(), 42);
+}
+
 auto main(int argc, char** argv) -> int {
+    MiniExecutor exec;
     ::testing::InitGoogleTest(&argc, argv);
     return RUN_ALL_TESTS();
 }