// -*- C++ -*-
// C++23 overlay: <generator>  (P2502R2)
// Provides std::generator<T> — a synchronous coroutine generator.

#ifndef _OMNI_CPP23_GENERATOR_OVERLAY
#define _OMNI_CPP23_GENERATOR_OVERLAY

#include <coroutine>
#include <exception>
#include <iterator>
#include <memory>
#include <type_traits>
#include <utility>

#if __cplusplus >= 202100L

#ifndef __cpp_lib_generator
#define __cpp_lib_generator 202207L

namespace std {

template <class Ref, class V = void, class Allocator = void>
class generator {
public:
  using value_type = conditional_t<is_void_v<V>, remove_cvref_t<Ref>, V>;
  using reference = conditional_t<is_void_v<V>, value_type&&, Ref>;
  using yielded = conditional_t<is_reference_v<reference>, reference, const reference&>;

  struct promise_type {
    using value_type = generator::value_type;

    const std::remove_reference_t<yielded>* value_ = nullptr;
    std::exception_ptr exception_ = nullptr;

    generator get_return_object() {
      return generator{coroutine_handle<promise_type>::from_promise(*this)};
    }

    suspend_always initial_suspend() noexcept { return {}; }
    suspend_always final_suspend() noexcept { return {}; }

    suspend_always yield_value(const value_type& val) noexcept {
      value_ = std::addressof(val);
      return {};
    }

    suspend_always yield_value(value_type&& val) noexcept {
      value_ = std::addressof(val);
      return {};
    }

    void return_void() noexcept {}

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

    // Disallow co_await in generator coroutines
    void await_transform() = delete;
  };

  class iterator {
  public:
    using value_type = generator::value_type;
    using difference_type = ptrdiff_t;
    using iterator_concept = input_iterator_tag;

  private:
    coroutine_handle<promise_type> handle_ = nullptr;

  public:
    iterator() = default;
    explicit iterator(coroutine_handle<promise_type> h) noexcept : handle_(h) {}

    const value_type& operator*() const {
      return *handle_.promise().value_;
    }

    iterator& operator++() {
      handle_.resume();
      if (handle_.done()) {
        auto& p = handle_.promise();
        if (p.exception_)
          std::rethrow_exception(p.exception_);
      }
      return *this;
    }

    void operator++(int) { ++*this; }

    friend bool operator==(const iterator& it, default_sentinel_t) {
      return it.handle_.done();
    }
  };

  // Constructors
  generator() = default;

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

  generator& operator=(generator&& other) noexcept {
    if (this != &other) {
      if (handle_) handle_.destroy();
      handle_ = std::exchange(other.handle_, nullptr);
    }
    return *this;
  }

  ~generator() {
    if (handle_) handle_.destroy();
  }

  generator(const generator&) = delete;
  generator& operator=(const generator&) = delete;

  iterator begin() {
    handle_.resume();
    if (handle_.done()) {
      auto& p = handle_.promise();
      if (p.exception_)
        std::rethrow_exception(p.exception_);
    }
    return iterator{handle_};
  }

  default_sentinel_t end() const noexcept { return {}; }

private:
  explicit generator(coroutine_handle<promise_type> h) noexcept : handle_(h) {}

  coroutine_handle<promise_type> handle_ = nullptr;
};

} // namespace std

#endif // __cpp_lib_generator
#endif // __cplusplus >= 202100L
#endif // _OMNI_CPP23_GENERATOR_OVERLAY
