feat: add then<>(..., []() -> Awaitable {})

include/rua/sync/then.hpp

@@ -11,6 +11,40 @@
 
 namespace rua {
 
+template <typename R, typename Awaitable, typename Callback>
+inline enable_if_t<!is_awaitable<R>::value, future<R>>
+_then(Awaitable &&awaitable, Callback &&callback) {
+	auto aw = wrap_awaiter(std::forward<Awaitable>(awaitable));
+	if (aw->await_ready()) {
+		return expected_invoke(std::forward<Callback>(callback), [&]() {
+			return aw->await_resume();
+		});
+	}
+
+	struct ctx_t {
+		awaiter_wrapper<Awaitable &&> aw;
+		decay_t<Callback> cb;
+	};
+
+	auto prm = new newable_promise<R, ctx_t>(
+		ctx_t{std::move(aw), std::forward<Callback>(callback)});
+
+	if (await_suspend(*prm->extend().aw, [prm]() {
+			prm->fulfill(expected_invoke(prm->extend().cb, [prm]() {
+				return prm->extend().aw->await_resume();
+			}));
+		})) {
+		return future<R>(*prm);
+	}
+
+	auto r_exp = expected_invoke(
+		prm->extend().cb, [prm]() { return prm->extend().aw->await_resume(); });
+
+	prm->unuse();
+
+	return r_exp;
+}
+
 template <
 	typename Awaitable,
 	typename Callback,
@@ -19,37 +53,70 @@ template <
 		std::declval<
 			warp_expected_t<decay_t<await_result_t<Awaitable &&>>>>())),
 	typename R = unwarp_expected_t<ExpR>>
-inline future<R> then(Awaitable &&awaitable, Callback &&callback) {
+inline decltype(_then<R>(
+	std::declval<Awaitable &&>(), std::declval<Callback &&>()))
+then(Awaitable &&awaitable, Callback &&callback) {
+	return _then<R>(
+		std::forward<Awaitable>(awaitable), std::forward<Callback>(callback));
+}
+
+template <
+	typename R,
+	typename Awaitable,
+	typename Callback,
+	typename R2 = unwarp_expected_t<decay_t<await_result_t<R>>>>
+inline enable_if_t<is_awaitable<R>::value, future<R2>>
+_then(Awaitable &&awaitable, Callback &&callback) {
 	auto aw = wrap_awaiter(std::forward<Awaitable>(awaitable));
 	if (aw->await_ready()) {
-		return expected_invoke(std::forward<Callback>(callback), [&]() {
+		auto r_exp = expected_invoke(std::forward<Callback>(callback), [&]() {
 			return aw->await_resume();
 		});
+
+		if (!r_exp) {
+			return r_exp.error();
+		}
+
+		return _then<R2>(*std::move(r_exp), [](expected<R2> r2_exp) {
+			return std::move(r2_exp);
+		});
 	}
 
 	struct ctx_t {
 		awaiter_wrapper<Awaitable &&> aw;
 		decay_t<Callback> cb;
 	};
 
-	auto prm = new newable_promise<R, ctx_t>(
+	auto prm = new newable_promise<R2, ctx_t>(
 		ctx_t{std::move(aw), std::forward<Callback>(callback)});
 
 	if (await_suspend(*prm->extend().aw, [prm]() {
-			prm->fulfill(expected_invoke(prm->extend().cb, [prm]() {
+			auto r_exp = expected_invoke(prm->extend().cb, [prm]() {
 				return prm->extend().aw->await_resume();
-			}));
+			});
+			if (!r_exp) {
+				prm->fulfill(r_exp.error());
+				return;
+			}
+			then(*std::move(r_exp), [prm](expected<R2> r2_exp) {
+				prm->fulfill(std::move(r2_exp));
+			});
 		})) {
-		return future<R>(*prm);
+		return future<R2>(*prm);
 	}
 
-	auto exp = expected_invoke(prm->extend().cb, [prm]() mutable {
-		return prm->extend().aw->await_resume();
-	});
+	auto r_exp = expected_invoke(
+		prm->extend().cb, [prm]() { return prm->extend().aw->await_resume(); });
 
 	prm->unuse();
 
-	return exp;
+	if (!r_exp) {
+		return r_exp.error();
+	}
+
+	return _then<R2>(*std::move(r_exp), [](expected<R2> r2_exp) {
+		return std::move(r2_exp);
+	});
 }
 
 namespace await_operators {