///////////////////////////////////////////////////////////////////////////////
// Reference implementation of std::generator proposal P2168.
//
// See https://wg21.link/P2168 for details.
//
///////////////////////////////////////////////////////////////////////////////
// Copyright Lewis Baker, Corentin Jabot
//
// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0.
// (See accompanying file LICENSE or http://www.boost.org/LICENSE_1_0.txt)
///////////////////////////////////////////////////////////////////////////////

export module nihil.generator:generator_promise_base;

import nihil.std;
import :elements_of;
import :forward;
import :manual_lifetime;

namespace nihil {

template <typename Ref>
struct generator_promise_base
{
	template <typename Ref2, typename Value, typename Alloc>
	friend class generator;

	generator_promise_base *m_root;
	std::coroutine_handle<> m_parent_or_leaf;

	// Note: Using manual_lifetime here to avoid extra calls to exception_ptr
	// constructor/destructor in cases where it is not needed (i.e. where this
	// generator coroutine is not used as a nested coroutine).
	// This member is lazily constructed by the __yield_sequence_awaiter::await_suspend()
	// method if this generator is used as a nested generator.
	manual_lifetime<std::exception_ptr> m_exception;
	manual_lifetime<Ref> m_value;

	explicit generator_promise_base(std::coroutine_handle<> this_coro) noexcept
		: m_root(this)
		, m_parent_or_leaf(this_coro)
	{
	}

	~generator_promise_base()
	{
		if (m_root != this) {
			// This coroutine was used as a nested generator and so will
			// have constructed its __exception_ member which needs to be
			// destroyed here.
			m_exception.destruct();
		}
	}

	auto initial_suspend() noexcept -> std::suspend_always
	{
		return {};
	}

	auto return_void() noexcept -> void
	{
	}

	auto unhandled_exception() -> void
	{
		if (m_root != this)
			m_exception.get() = std::current_exception();
		else
			throw;
	}

	// Transfers control back to the parent of a nested coroutine
	struct final_awaiter
	{
		auto await_ready() noexcept -> bool
		{
			return false;
		}

		template <typename Promise>
		auto
		await_suspend(std::coroutine_handle<Promise> h) noexcept -> std::coroutine_handle<>
		{
			Promise &promise = h.promise();
			generator_promise_base &root = *promise.m_root;
			if (&root != &promise) {
				auto parent = promise.m_parent_or_leaf;
				root.m_parent_or_leaf = parent;
				return parent;
			}
			return std::noop_coroutine();
		}

		void await_resume() noexcept
		{
		}
	};

	auto final_suspend() noexcept -> final_awaiter
	{
		return {};
	}

	auto yield_value(Ref &&x) noexcept(std::is_nothrow_move_constructible_v<Ref>)
		-> std::suspend_always
	{
		m_root->m_value.construct((Ref &&)x);
		return {};
	}

	template <typename T>
	requires(!std::is_reference_v<Ref>) && std::is_convertible_v<T, Ref>
	auto
	yield_value(T &&x) noexcept(std::is_nothrow_constructible_v<Ref, T>) -> std::suspend_always
	{
		m_root->m_value.construct((T &&)x);
		return {};
	}

	template <typename Gen>
	struct yield_sequence_awaiter
	{
		Gen m_gen;

		yield_sequence_awaiter(Gen &&g) noexcept
			// Taking ownership of the generator ensures frame are destroyed
		        // in the reverse order of their execution.
			: m_gen((Gen &&)g)
		{
		}

		auto await_ready() noexcept -> bool
		{
			return false;
		}

		// set the parent, root and exceptions pointer and
		// resume the nested
		template <typename Promise>
		auto
		await_suspend(std::coroutine_handle<Promise> h) noexcept -> std::coroutine_handle<>
		{
			generator_promise_base &current = h.promise();
			generator_promise_base &nested = *m_gen.get_promise();
			generator_promise_base &root = *current.m_root;

			nested.m_root = current.m_root;
			nested.m_parent_or_leaf = h;

			// Lazily construct the __exception_ member here now that we
			// know it will be used as a nested generator. This will be
			// destroyed by the promise destructor.
			nested.m_exception.construct();
			root.m_parent_or_leaf = m_gen.get_coro();

			// Immediately resume the nested coroutine (nested generator)
			return m_gen.get_coro();
		}

		void await_resume()
		{
			generator_promise_base &nested_promise = *m_gen.get_promise();

			if (nested_promise.m_exception.get()) {
				std::rethrow_exception(nested_promise.m_exception.get());
			}
		}
	};

	template <typename OValue, typename OAlloc>
	auto yield_value(nihil::elements_of<generator<Ref, OValue, OAlloc>> g) noexcept
		-> yield_sequence_awaiter<generator<Ref, OValue, OAlloc>>

	{
		return std::move(g).get();
	}

	template <std::ranges::range Rng, typename Allocator>
	auto yield_value(nihil::elements_of<Rng, Allocator> &&x)
		-> yield_sequence_awaiter<generator<Ref, std::remove_cvref_t<Ref>, Allocator>>

	{
		return [](std::allocator_arg_t, Allocator,
		          auto &&rng) -> generator<Ref, std::remove_cvref_t<Ref>, Allocator> {
			for (auto &&e : rng)
				co_yield static_cast<decltype(e)>(e);
		}(std::allocator_arg, x.get_allocator(), std::forward<Rng>(x.get()));
	}

	void resume()
	{
		m_parent_or_leaf.resume();
	}

	// Disable use of co_await within this coroutine.
	void await_transform() = delete;
};

} // namespace nihil