Fixed tail-call implementation for lazy_task amd when_all.

If await_suspend() returns a std::coroutine_handle then
this is unconditionally resumed. So we now return a
dummy coroutine_handle whose resume() method just
immediately suspends in cases where there is no
continuation ready to be resumed.

Author: lewissbaker

include/cppcoro/detail/continuation.hpp

@@ -15,7 +15,7 @@ namespace cppcoro
 		{
 		public:
 
-			using callback_t = void(void*);
+			using callback_t = std::experimental::coroutine_handle<>(void*);
 
 			continuation() noexcept
 				: m_callback(nullptr)
@@ -39,26 +39,18 @@ namespace cppcoro
 
 			void resume() noexcept
 			{
-				if (m_callback == nullptr)
-				{
-					std::experimental::coroutine_handle<>::from_address(m_state).resume();
-				}
-				else
-				{
-					m_callback(m_state);
-				}
+				tail_call_resume().resume();
 			}
 
-			std::experimental::coroutine_handle<> resume_or_get_coroutine_handle()
+			std::experimental::coroutine_handle<> tail_call_resume()
 			{
 				if (m_callback == nullptr)
 				{
 					return std::experimental::coroutine_handle<>::from_address(m_state);
 				}
 				else
 				{
-					m_callback(m_state);
-					return {};
+					return m_callback(m_state);
 				}
 			}
 

include/cppcoro/detail/dummy_coroutine.hpp

@@ -0,0 +1,61 @@
+///////////////////////////////////////////////////////////////////////////////
+// Copyright (c) Lewis Baker
+// Licenced under MIT license. See LICENSE.txt for details.
+///////////////////////////////////////////////////////////////////////////////
+#ifndef CPPCORO_DETAIL_DUMMY_COROUTINE_HPP_INCLUDED
+#define CPPCORO_DETAIL_DUMMY_COROUTINE_HPP_INCLUDED
+
+#include <experimental/coroutine>
+
+namespace cppcoro
+{
+	namespace detail
+	{
+		class dummy_coroutine final
+		{
+		public:
+
+			struct promise_type final
+			{
+				dummy_coroutine get_return_object() noexcept;
+				std::experimental::suspend_never initial_suspend() noexcept { return {}; }
+				std::experimental::suspend_never final_suspend() noexcept { return {}; }
+				void unhandled_exception() noexcept {}
+				void return_void() noexcept {}
+			};
+
+			static std::experimental::coroutine_handle<> handle() noexcept
+			{
+				static const auto dummyHandle = coroutine().m_handle;
+				return dummyHandle;
+			}
+
+		private:
+
+
+			explicit dummy_coroutine(std::experimental::coroutine_handle<> handle) noexcept
+				: m_handle(handle)
+			{}
+
+			static dummy_coroutine coroutine()
+			{
+				for (;;)
+				{
+					co_await std::experimental::suspend_always{};
+				}
+			}
+
+			const std::experimental::coroutine_handle<> m_handle;
+
+		};
+
+		inline dummy_coroutine dummy_coroutine::promise_type::get_return_object() noexcept
+		{
+			return dummy_coroutine{
+				std::experimental::coroutine_handle<promise_type>::from_promise(*this)
+			};
+		}
+	}
+}
+
+#endif

include/cppcoro/detail/when_all_awaitable.hpp

@@ -6,6 +6,7 @@
 #define CPPCORO_DETAIL_WHEN_ALL_AWAITABLE_HPP_INCLUDED
 
 #include <cppcoro/detail/continuation.hpp>
+#include <cppcoro/detail/dummy_coroutine.hpp>
 
 #include <atomic>
 
@@ -17,7 +18,7 @@ namespace cppcoro
 		{
 		public:
 
-			when_all_awaitable(std::size_t count) noexcept
+			explicit when_all_awaitable(std::size_t count) noexcept
 				: m_refCount(count + 1)
 			{}
 
@@ -31,22 +32,33 @@ namespace cppcoro
 				return m_refCount.load(std::memory_order_acquire) == 1;
 			}
 
-			bool await_suspend(std::experimental::coroutine_handle<> awaiter) noexcept
+			auto await_suspend(std::experimental::coroutine_handle<> awaiter) noexcept
 			{
 				m_awaiter = awaiter;
-				return m_refCount.fetch_sub(1, std::memory_order_acq_rel) > 1;
+				if (m_refCount.fetch_sub(1, std::memory_order_acq_rel) == 1)
+				{
+					return awaiter;
+				}
+				else
+				{
+					return dummy_coroutine::handle();
+				}
 			}
 
 			void await_resume() noexcept {}
 
 		private:
 
-			static void resumer_callback(void* state) noexcept
+			static std::experimental::coroutine_handle<> resumer_callback(void* state) noexcept
 			{
 				auto* that = static_cast<when_all_awaitable*>(state);
 				if (that->m_refCount.fetch_sub(1, std::memory_order_acq_rel) == 1)
 				{
-					that->m_awaiter.resume();
+					return that->m_awaiter;
+				}
+				else
+				{
+					return dummy_coroutine::handle();
 				}
 			}
 

include/cppcoro/lazy_task.hpp

@@ -51,7 +51,7 @@ namespace cppcoro
 					std::experimental::coroutine_handle<>
 					await_suspend([[maybe_unused]] std::experimental::coroutine_handle<> coroutine)
 					{
-						return m_continuation.resume_or_get_coroutine_handle();
+						return m_continuation.tail_call_resume();
 					}
 
 					void await_resume() noexcept {}

lib/build.cake

@@ -42,6 +42,7 @@ includes = cake.path.join(env.expand('${CPPCORO}'), 'include', 'cppcoro', [
 
 detailIncludes = cake.path.join(env.expand('${CPPCORO}'), 'include', 'cppcoro', 'detail', [
   'continuation.hpp',
+  'dummy_coroutine.hpp',
   'when_all_awaitable.hpp',
   'unwrap_reference.hpp',
   ])