// generator - Single-header, ranges-compatible generator type built
// on C++20 coroutines
// Written in 2021 by Sy Brand (tartanllama@gmail.com, @TartanLlama)
//
// Documentation available at https://tl.tartanllama.xyz/
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to the
// public domain worldwide. This software is distributed without any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software. If not, see
// <http://creativecommons.org/publicdomain/zero/1.0/>.
export module nihil.core:generator;

import nihil.std;

namespace nihil {

export template <typename T>
struct generator
{
private:
	struct promise
	{
		using value_type = std::remove_reference_t<T>;
		using reference_type =
			std::conditional_t<std::is_pointer_v<value_type>, value_type, value_type &>;
		using pointer_type =
			std::conditional_t<std::is_pointer_v<value_type>, value_type, value_type *>;

		promise() = default;

		[[nodiscard]] auto get_return_object() -> generator
		{
			return generator(std::coroutine_handle<promise>::from_promise(*this));
		}

		[[nodiscard]] auto initial_suspend() const -> std::suspend_always
		{
			return {};
		}

		[[nodiscard]] auto final_suspend() const noexcept -> std::suspend_always
		{
			return {};
		}

		void return_void() const noexcept
		{
		}

		void unhandled_exception() noexcept
		{
			exception_ = std::current_exception();
		}

		void rethrow_if_exception()
		{
			if (exception_) {
				std::rethrow_exception(exception_);
			}
		}

		std::suspend_always yield_value(value_type &&v) noexcept
		{
			value_ = std::move(v);
			return {};
		}

		std::suspend_always yield_value(value_type const &v) noexcept
		requires(!std::is_reference_v<T>)
		{
			value_ = v;
			return {};
		}

		std::suspend_always yield_value(value_type &v) noexcept
		requires(std::is_reference_v<T>)
		{
			value_ = &v;
			return {};
		}

		std::exception_ptr                                     exception_;
		std::variant<std::monostate, value_type, value_type *> value_;
	};

public:
	using promise_type = promise;
	class sentinel
	{
	};

	class iterator
	{
		using handle_type = std::coroutine_handle<promise_type>;

	public:
		using value_type = typename promise_type::value_type;
		using reference_type = typename promise_type::reference_type;
		using pointer_type = typename promise_type::pointer_type;
		using difference_type = std::ptrdiff_t;

		iterator() = default;
		~iterator()
		{
			if (handle_)
				handle_.destroy();
		}

		// Non-copyable because coroutine handles point to a unique resource
		iterator(iterator const &) = delete;
		iterator(iterator &&rhs) noexcept
			: handle_(std::exchange(rhs.handle_, nullptr))
		{
		}

		auto operator=(iterator const &) -> iterator & = delete;

		auto operator=(iterator &&rhs) noexcept -> iterator &
		{
			handle_ = std::exchange(rhs.handle_, nullptr);
			return *this;
		}

		friend auto operator==(iterator const &it, sentinel) noexcept -> bool
		{
			return (!it.handle_ || it.handle_.done());
		}

		auto operator++() -> iterator &
		{
			handle_.resume();
			if (handle_.done()) {
				handle_.promise().rethrow_if_exception();
			}
			return *this;
		}

		void operator++(int)
		{
			(void)this->operator++();
		}

		reference_type operator*() const
			noexcept(noexcept(std::is_nothrow_copy_constructible_v<reference_type>))
		requires(std::is_reference_v<T>)
		{
			return *std::get<value_type *>(handle_.promise().value_);
		}

		reference_type operator*() const
			noexcept(noexcept(std::is_nothrow_copy_constructible_v<reference_type>))
		requires(!std::is_reference_v<T>)
		{
			return std::get<value_type>(handle_.promise().value_);
		}

		value_type *operator->() const
			noexcept(noexcept(std::is_nothrow_copy_constructible_v<reference_type>))
		requires(std::is_reference_v<T>)
		{
			return std::get<value_type *>(handle_.promise().value_);
		}

		value_type *operator->() const
			noexcept(noexcept(std::is_nothrow_copy_constructible_v<reference_type>))
		requires(!std::is_reference_v<T>)
		{
			return &std::get<value_type>(handle_.promise().value_);
		}

	private:
		friend struct generator;

		explicit iterator(handle_type handle)
			: handle_(handle)
		{
		}

		handle_type handle_;
	};

	using handle_type = std::coroutine_handle<promise_type>;

	generator() noexcept = default;
	~generator()
	{
		if (handle_)
			handle_.destroy();
	}

	generator(generator const &) = delete;

	generator(generator &&rhs) noexcept
		: handle_(std::exchange(rhs.handle_, nullptr))
	{
	}

	auto operator=(generator const &) -> generator & = delete;

	auto operator=(generator &&rhs) noexcept -> generator &
	{
		swap(rhs);
		return *this;
	}

	auto begin() -> iterator
	{
		handle_.resume();
		if (handle_.done()) {
			handle_.promise().rethrow_if_exception();
		}
		return iterator(std::exchange(handle_, nullptr));
	}

	auto end() const noexcept -> sentinel
	{
		return {};
	}

	void swap(generator &other) noexcept
	{
		std::swap(handle_, other.handle_);
	}

private:
	friend class iterator;
	explicit generator(handle_type handle) noexcept
		: handle_(handle)
	{
	}

	handle_type handle_ = nullptr;
};
} // namespace nihil

export template <typename T>
inline constexpr bool std::ranges::enable_view<nihil::generator<T>> = true;