Example of a Python-like generator implemented using C++20 coroutines

generator.hpp

/**
 * @file generator.hpp
 * @author andrea.mecchia@mail.polimi.it
 * @brief Example of a Python-like generator implemented
 * using C++20 coroutines
 * @date 2022-04-10
 * 
 * @copyright Copyright (c) 2022 Andrea Mecchia
 * 
 * C++20 coroutines allow us to create Python-like
 * generators in C++.
 * 
 * The following is a basic example of how to implement a
 * generic generator type.
 * 
 * Usage of the generator:
 * 
 * ```cpp
 * Generator<int> range(int start, int end, int step = 1)
 * {
 *     for (int i = start; i < end; i += step)
 *     {
 *         co_yield i;
 *     }
 * }
 * 
 * int main()
 * {
 *     for (auto i : range(10, 100, 10))
 *     {
 *         printf("%d\n", i);
 *     }
 * 
 * 	   // Or
 *     auto gen = range(10, 100, 10);
 *     while (gen())
 *     {
 *         printf("%d\n", *gen);
 *     }
 * }
 * ```
 * 
 * See https://en.cppreference.com/w/cpp/language/coroutines
 */

#pragma once

#include <assert.h>

#include <iterator>
#include <optional>
#include <coroutine>


template<typename> class Generator;


/**
 * @brief Iterator type used to iterate the values
 * yielded by a generator object.
 *
 * @tparam T the type of the yield values
 */
template<typename T>
class GeneratorIterator
{
	using SelfT = GeneratorIterator;
	using GeneratorT = Generator<T>;
	using RefT = T&;
	using PtrT = T*;

public:
	// ================================
	// Iterator traits
	// ================================

	using iterator_category = std::forward_iterator_tag;
	using difference_type   = void;
	using value_type        = T;
	using reference         = RefT;
	using pointer           = PtrT;

	/**
	 * @brief Construct a new iterator for the given
	 * generator.
	 *
	 * @param gen_in ref to the generator to iterate
	 */
	GeneratorIterator(GeneratorT& gen_in)
		: gen{gen_in}
	{
		// Coroutine starts in suspend state, resume here
		gen();
	}

	/**
	 * @brief Returns a ref to the last yielded value.
	 */
	inline RefT operator*()
	{
		return *gen;
	}

	/**
	 * @brief Returns a ptr to the last yielded value.
	 */
	PtrT operator->()
	{
		return &(**this);
	}

	/**
	 * @brief Returns true only if this iterator and the
	 * given iterator refer to the same coroutine.
	 *
	 * @param other another gen iterator
	 * @return true if they refer to the same coroutine
	 * @return false otherwise
	 */
	inline bool operator==(SelfT const& other)
	{
		// True if they refer to the same coroutine
		return gen == other.gen;
	}

	/**
	 * @brief Returns true only if the coroutine has
	 * reached the end of execution.
	 */
	bool operator==(nullptr_t)
	{
		// The end iterator is a null pointer, so this is true if the coroutine has reached the end
		return gen.done();
	}

	/**
	 * @brief Returns false if this iterator and the
	 * given iterator refer to different coroutines.
	 *
	 * @param other another gen iterator
	 * @return true if they refer to different
	 * coroutines
	 * @return false otherwise
	 */
	bool operator!=(SelfT const& other)
	{
		return !(*this == other);
	}

	/**
	 * @brief Returns true if coroutine has not reached
	 * the end of execution.
	 */
	inline bool operator!=(nullptr_t)
	{
		return !(*this == nullptr);
	}

	/**
	 * @brief Resume coroutine, yield next value. Return
	 * ref to self.
	 */
	SelfT& operator++()
	{
		// Resume coroutine
		return gen(), *this;
	}

private:
	/* The generator object. */
	GeneratorT& gen;
};


/**
 * @brief Wrapper for coroutines that generate values of
 * the given type.
 *
 * The coroutine starts in a suspend state, before
 * reading the next value it must be resume by calling
 * the generator or using `next()`.
 *
 * Resuming the coroutine after it reached the end of
 * execution has undefined behavior. Use `done()` or
 * cast to `bool` before resuming after the first time.
 *
 * After the generator has finished, you can still read
 * the last yielded value.
 *
 * @tparam T value of the yield type
 */
template<typename T>
class Generator
{
	class Promise
	{
		friend Generator;

		using HandleT = std::coroutine_handle<Promise>;

	public:
		/* Return the generator object. */
		Generator get_return_object()
		{
			return {HandleT::from_promise(*this)};
		}

		/* Always suspend at the beginning. */
		constexpr std::suspend_always initial_suspend()
		{
			return {};
		}

		/* Never suspend at the end. */
		constexpr std::suspend_never final_suspend() noexcept
		{
			return {};
		}

		/* Called when the routine does not handle the
			exception */
		void unhandled_exception() {}

		/* Called when the routine returns void. */
		void return_void() {}

		/* Store the yield value. */
		std::suspend_always yield_value(auto&& ...createArgs)
		{
			// Set the valut to yield
			value = T{std::forward<decltype(createArgs)>(createArgs)...};
			return {};
		}

	private:
		/* The next value to return.
			Optional prevents constructing the object. */
		std::optional<T> value;
	};

public:
	using promise_type = Promise;

	using YieldT    = T;
	using IteratorT = GeneratorIterator<T>;

	/**
	 * @brief Construct a new Generator object from the
	 * coroutine handle.
	 *
	 * @param handle_in the coroutine handle
	 */
	Generator(Promise::HandleT&& handle_in)
		: handle{std::move(handle_in)}
	{}

	/**
	 * @brief Resume the coroutine.
	 *
	 * @return true if coroutine has reached the end of
	 * execution
	 * @return false otherwise
	 */
	bool operator()()
	{
		assert(!!handle);

		// Resume coroutine
		return handle(), !handle.done();
	}

	/**
	 * @brief Resume the coroutine and return the
	 * next generated value.
	 *
	 * @return next generated value
	 */
	YieldT& next()
	{
		assert(!!handle);

		// Coroutine starts in suspend state, we resume here
		handle();

		// Return current yield value
		return *handle.promise().value;
	}

	/**
	 * @brief Returns a ref to the last yielded value.
	 * @{
	 */
	inline YieldT& operator*()
	{
		return *handle.promise().value;
	}

	inline YieldT const& operator*() const
	{
		return *const_cast<Generator&>(*this);
	}
	/** @} */

	/**
	 * @brief Returns true if there are no more values
	 * to generate.
	 * @{
	 */
	inline bool done() const
	{
		assert(!!handle);
		return handle.done();
	}

	inline operator bool() const
	{
		assert(!!handle);
		return handle.done();
	}
	/** @} */

	/**
	 * @brief Returns an iterator for this generator.
	 * 
	 * Note that this only points to the beginning of
	 * the generator if the generator was never resumed.
	 */
	inline IteratorT begin()
	{
		return {*this};
	}

	/**
	 * @brief Returns an iterator pointing to the end
	 * of this generator.
	 */
	constexpr auto end()
	{
		return nullptr;
	}

protected:
	/* The coroutine handle. */
	Promise::HandleT handle;
};