From 1bfbaaa840410e2277cd953fa36c4e8e225949cb Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Thu, 4 Mar 2021 18:33:58 +0300
Subject: [PATCH] Support async::future::then

---
 libs/async/include/psemek/async/future.hpp | 121 +++++++++++++++++++--
 1 file changed, 112 insertions(+), 9 deletions(-)

diff --git a/libs/async/include/psemek/async/future.hpp b/libs/async/include/psemek/async/future.hpp
index 782ecfb3..16fd89b4 100644
--- a/libs/async/include/psemek/async/future.hpp
+++ b/libs/async/include/psemek/async/future.hpp
@@ -38,6 +38,18 @@ namespace psemek::async
 			using type = bool;
 		};
 
+		template <typename T>
+		struct then_function
+		{
+			using type = util::function<void(T &)>;
+		};
+
+		template <>
+		struct then_function<void>
+		{
+			using type = util::function<void()>;
+		};
+
 		template <typename T>
 		struct task_state
 		{
@@ -50,6 +62,9 @@ namespace psemek::async
 
 			std::condition_variable value_cv;
 
+			std::mutex then_mutex;
+			typename then_function<T>::type then_func;
+
 			task_state(bool auto_cancel)
 				: auto_cancel{auto_cancel}
 			{}
@@ -159,6 +174,9 @@ namespace psemek::async
 			return wait_for(std::chrono::seconds{0});
 		}
 
+		template <typename F>
+		auto then(F && f);
+
 	private:
 		std::shared_ptr<detail::task_state<T>> state_;
 
@@ -202,17 +220,25 @@ namespace psemek::async
 		void set_value(T const & value)
 		{
 			if (!state_) throw empty_promise_error{};
-			std::lock_guard lock{state_->value_mutex};
-			if (state_->value || state_->exception) throw satisfied_promise_error{};
-			state_->value = value;
+			{
+				std::lock_guard lock{state_->value_mutex};
+				if (state_->value || state_->exception) throw satisfied_promise_error{};
+				state_->value = value;
+			}
+			if (state_->then_func)
+				state_->then_func(*state_->value);
 		}
 
 		void set_value(T && value)
 		{
 			if (!state_) throw empty_promise_error{};
-			std::lock_guard lock{state_->value_mutex};
-			if (state_->value || state_->exception) throw satisfied_promise_error{};
-			state_->value = std::move(value);
+			{
+				std::lock_guard lock{state_->value_mutex};
+				if (state_->value || state_->exception) throw satisfied_promise_error{};
+				state_->value = std::move(value);
+			}
+			if (state_->then_func)
+				state_->then_func(*state_->value);
 		}
 
 		void set_exception(std::exception_ptr e)
@@ -254,9 +280,13 @@ namespace psemek::async
 		void set_value()
 		{
 			if (!state_) throw empty_promise_error{};
-			std::lock_guard lock{state_->value_mutex};
-			if (state_->value || state_->exception) throw satisfied_promise_error{};
-			state_->value = true;
+			{
+				std::lock_guard lock{state_->value_mutex};
+				if (state_->value || state_->exception) throw satisfied_promise_error{};
+				state_->value = true;
+			}
+			if (state_->then_func)
+				state_->then_func();
 		}
 
 		void set_exception(std::exception_ptr e)
@@ -333,4 +363,77 @@ namespace psemek::async
 		util::function<R(Args...)> func_;
 	};
 
+	template <typename T>
+	template <typename F>
+	auto future<T>::then(F && f)
+	{
+		if (!state_) throw empty_future_error{};
+
+		std::lock_guard lock{state_->value_mutex};
+		if (state_->value || state_->exception)
+		{
+			if constexpr (std::is_same_v<T, void>)
+			{
+				using R = decltype(f());
+				promise<R> p;
+
+				if (state_->value)
+				{
+					if constexpr (std::is_same_v<R, void>)
+					{
+						std::forward<F>(f)();
+						p.set_value();
+					}
+					else
+					{
+						p.set_value(std::forward<F>(f)());
+					}
+				}
+				else
+					p.set_exception(state_->exception);
+				return p.get_future();
+			}
+			else
+			{
+				using R = decltype(f(*(state_->value)));
+				promise<R> p;
+
+				if (state_->value)
+				{
+					if constexpr (std::is_same_v<R, void>)
+					{
+						std::forward<F>(f)(*(state_->value));
+						p.set_value();
+					}
+					else
+					{
+						p.set_value(std::forward<F>(f)(*(state_->value)));
+					}
+				}
+				else
+					p.set_exception(state_->exception);
+				return p.get_future();
+			}
+		}
+
+		if constexpr (std::is_same_v<T, void>)
+		{
+			using R = decltype(f());
+			packaged_task<R()> t(std::forward<F>(f));
+
+			std::lock_guard lock{state_->then_mutex};
+			state_->then_func = std::move(t);
+			return t.get_future();
+		}
+		else
+		{
+			using R = decltype(f(*(state_->value)));
+			packaged_task<R(T &)> t(std::forward<F>(f));
+
+			std::lock_guard lock{state_->then_mutex};
+			state_->then_func = std::move(t);
+			return t.get_future();
+		}
+	}
+
 }