Closed danakj closed 1 year ago
A rough sketch. Without any integration into the Iterator concept and IteratorImpl class. But I think this can subclass from IteratorImpl and it Just Works(tm). Then any function can return an IterGenerator<T>
and yield T
s and it's a sus::Iterator
.
Some questions:
AsyncIterator
or something, which is bool is_done()
or something. This would allow the iterator to report separately that it's done iterating vs it returned a None. Sync iterators can't have this distinction. But is this useful?co_await
and return control up. How does that fit into this model and make it customizable?https://godbolt.org/z/Ejhq7s8nx
#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 << " ";
}
}
await in an iterator does not seem particularly useful. it would need to yield and then have the caller give it data that it can await for...?
Iterators can yield other iterators with for loops:
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);
}
}
Prints: 1 2 4 7 10 9 8
Delete await_transform()
in the promise_type
to avoid co_await in the IteratorGenerator since it's a sync iterator and next() needs to return the next value. Returning early won't work, the caller expects a value there, there's no poll API.
void await_transform() = delete;
Subspace, like Rust, currently lacks the control-flow "register when iterating: https://without.boats/blog/the-registers-of-rust/
Maybe we would like an AsyncIterator that can be returned from a coroutine function?
Time to go learn about coroutines in depth.
See also prior art, all the methods on IEnumerable in C# .