Iterating + Control Flow about subspace HOT 3 CLOSED

#include <cassert>
#include <concepts>
#include <coroutine>
#include <cstdlib>
#include <iostream>
#include <utility>

namespace sus {
void check(bool v) { assert(v); }
[[noreturn]] void unreachable() { abort(); }
}  // namespace sus

template <class T>
struct Option {
    constexpr ~Option() noexcept {
        if (has) t_.~T();
    }
    union {
        T t_;
    };
    bool has = false;

    constexpr Option() = default;
    constexpr Option(T r) : t_(std::move(r)), has(true) {}

    Option(Option&& o) : has(std::exchange(o.has, false)) {
        if (has) {
            if constexpr (std::is_trivially_default_constructible_v<T>)
                t_ = o.t_;
            else
                new (&t_) T(std::move(o.t_));
            o.t_.~T();
        }
    }
    Option& operator=(Option&& o) {
        if (has) {
            t_.~T();
        }
        has = std::exchange(o.has, false);
        if (has) {
            if constexpr (std::is_trivially_default_constructible_v<T>)
                t_ = o.t_;
            else
                new (&t_) T(std::move(o.t_));
            o.t_.~T();
        }
        return *this;
    }

    constexpr void insert(T r) {
        if (!has) {
            has = true;
            new (&t_) T(std::move(r));
        } else {
            t_ = std::move(r);
        }
    }

    constexpr T unwrap_unchecked() && noexcept {
        has = false;
        T r = std::move(t_);
        t_.~T();
        return r;
    }

    constexpr Option take() & noexcept {
        if (has) {
            auto o = Option(std::move(t_));
            t_.~T();
            has = false;
            return o;
        } else {
            return Option();
        }
    }
};

template <class IG, class T>
struct IterPromise {
    IG get_return_object() { return IG(*this); }

    constexpr auto yield_value(Option<T> y) noexcept {
        yielded_ = std::move(y);
        return std::suspend_always();
    }

    constexpr auto initial_suspend() noexcept { return std::suspend_always(); }
    constexpr auto final_suspend() noexcept { return std::suspend_always(); }
    void unhandled_exception() noexcept { ::sus::unreachable(); }

    Option<T> yielded_;
};

struct IteratorEnd {};

template <class T>
struct IterGenerator {
    // Coroutine implementation.
    using promise_type = IterPromise<IterGenerator<T>, T>;

    using CoHandle = std::coroutine_handle<promise_type>;

    constexpr IterGenerator(promise_type& p) noexcept
        : co_handle_(CoHandle::from_promise(p)) {}
    constexpr ~IterGenerator() noexcept {
        if (co_handle_ != nullptr) co_handle_.destroy();
    }

    constexpr IterGenerator(IterGenerator&& o) noexcept
        : co_handle_(std::exchange(o.co_handle_, nullptr)) {
        ::sus::check(co_handle_ != nullptr);
    }
    constexpr IterGenerator& operator=(IterGenerator&& o) noexcept {
        co_handle_ = std::exchange(o.co_handle_, nullptr);
        ::sus::check(co_handle_ != nullptr);
    }

    constexpr Option<T> next() noexcept {
        if (!co_handle_.done()) co_handle_.resume();
        return co_handle_.promise().yielded_.take();
    }

    constexpr auto begin() & noexcept {
        // Ensure the first item is yielded and ready to be returned from the
        // iterator's operator*().
        if (!co_handle_.done()) co_handle_.resume();
        return IterGeneratorLoop(*this);
    }

    constexpr auto end() & noexcept { return IteratorEnd(); }

    struct IterGeneratorLoop {
        IterGeneratorLoop(IterGenerator& ig [[clang::lifetimebound]])
            : ig_(ig) {}
        constexpr bool operator==(const IteratorEnd&) noexcept {
            return ig_.co_handle_.done();
        }
        constexpr IterGeneratorLoop& operator++() & noexcept {
            ig_.co_handle_.resume();
            return *this;
        }
        constexpr T operator*() & noexcept {
            return ig_.co_handle_.promise().yielded_.take().unwrap_unchecked();
        }

        IterGenerator& ig_;
    };

    CoHandle co_handle_;
};

int main() {
    auto x = []() -> IterGenerator<int> {
        co_yield 1;
        co_yield 2;
        co_yield 4;
        co_yield 9;
    };
    auto it = x();

    std::cout << it.next().unwrap_unchecked() << "\n";

    for (int i : it) {
        std::cout << i << " ";
    }
}

int main() {
    auto y = []() -> IterGenerator<int> {
        co_yield 10;
        co_yield 9;
        co_yield 8;
    };
    auto x = [](IterGenerator<int> o) -> IterGenerator<int> {
        co_yield 1;
        co_yield 2;
        co_yield 4;
        co_yield 7;
        for (auto&& x : o) co_yield x;
    };
    auto it = x(y());

    print(it.next().unwrap_unchecked());

    for (int i : it) {
        print(i);
    }
}

    void await_transform() = delete;