Fix awaiter environment

include/beman/task/detail/awaiter.hpp

@@ -46,20 +46,23 @@ struct awaiter_op_t<Awaiter, ParentPromise, false> {
 };
 
 template <typename Value, typename Env, typename OwnPromise, typename ParentPromise>
-class awaiter : public ::beman::task::detail::state_base<Value, Env>,
-                ::beman::task::detail::state_rep<Env, ::beman::task::detail::handle<OwnPromise>> {
+class awaiter : public ::beman::task::detail::state_base<Value, Env> {
   public:
     using stop_token_type = typename ::beman::task::detail::state_base<Value, Env>::stop_token_type;
     using scheduler_type  = typename ::beman::task::detail::state_base<Value, Env>::scheduler_type;
 
-    explicit awaiter(::beman::task::detail::handle<OwnPromise> h)
-        : ::beman::task::detail::state_rep<Env, ::beman::task::detail::handle<OwnPromise>>(std::move(h)) {}
+    explicit awaiter(::beman::task::detail::handle<OwnPromise> h) : handle(std::move(h)) {}
     constexpr auto await_ready() const noexcept -> bool { return false; }
-    auto           await_suspend(::std::coroutine_handle<ParentPromise> parent) noexcept {
+    struct env_receiver {
+        ParentPromise* parent;
+        auto           get_env() const noexcept { return parent->get_env(); }
+    };
+    auto await_suspend(::std::coroutine_handle<ParentPromise> parent) noexcept {
+        this->state_rep.emplace(env_receiver{&parent.promise()});
         this->scheduler.emplace(
             this->template from_env<scheduler_type>(::beman::execution::get_env(parent.promise())));
         this->parent = ::std::move(parent);
-        return this->receiver.start(this);
+        return this->handle.start(this);
     }
     auto await_resume() { return this->result_resume(); }
 
@@ -82,20 +85,20 @@ class awaiter : public ::beman::task::detail::state_base<Value, Env>,
         return this->actual_complete();
     }
     auto actual_complete() -> std::coroutine_handle<> {
-        return this->::beman::task::detail::state_base<Value, Env>::no_completion_set()
-                   ? this->parent.promise().unhandled_stopped()
-                   : ::std::move(this->parent);
+        return this->no_completion_set() ? this->parent.promise().unhandled_stopped() : ::std::move(this->parent);
     }
     auto do_get_scheduler() -> scheduler_type override { return *this->scheduler; }
     auto do_set_scheduler(scheduler_type other) -> scheduler_type override {
         return ::std::exchange(*this->scheduler, other);
     }
     auto do_get_stop_token() -> stop_token_type override { return {}; }
-    auto do_get_environment() -> Env& override { return this->context; }
+    auto do_get_environment() -> Env& override { return this->state_rep->context; }
 
-    ::std::optional<scheduler_type>                       scheduler;
-    ::std::coroutine_handle<ParentPromise>                parent{};
-    ::std::optional<awaiter_op_t<awaiter, ParentPromise>> reschedule{};
+    ::beman::task::detail::handle<OwnPromise>                            handle;
+    ::std::optional<::beman::task::detail::state_rep<Env, env_receiver>> state_rep;
+    ::std::optional<scheduler_type>                                      scheduler;
+    ::std::coroutine_handle<ParentPromise>                               parent{};
+    ::std::optional<awaiter_op_t<awaiter, ParentPromise>>                reschedule{};
 };
 } // namespace beman::task::detail
 

include/beman/task/detail/handle.hpp

@@ -32,7 +32,11 @@ class handle {
     auto release() -> ::std::coroutine_handle<P> {
         return ::std::coroutine_handle<P>::from_promise(*this->h.release());
     }
-    auto get_env() const noexcept { return ::beman::execution::get_env(*this->h); }
+    P*   get() const noexcept { return this->h.get(); }
+    auto get_env() const noexcept {
+        assert(this->h.get());
+        return ::beman::execution::get_env(*this->h);
+    }
 };
 
 } // namespace beman::task::detail

include/beman/task/detail/promise_type.hpp

@@ -109,7 +109,11 @@ class promise_type
     auto get_scheduler() const noexcept -> scheduler_type { return this->get_state()->get_scheduler(); }
     auto get_allocator() const noexcept -> allocator_type { return this->allocator; }
     auto get_stop_token() const noexcept -> stop_token_type { return this->get_state()->get_stop_token(); }
-    auto get_environment() const noexcept -> const Environment& { return this->get_state()->get_environment(); }
+    auto get_environment() const noexcept -> const Environment& {
+        assert(this);
+        assert(this->get_state());
+        return this->get_state()->get_environment();
+    }
 
   private:
     using env_t = ::beman::task::detail::promise_env<promise_type>;

include/beman/task/detail/state_base.hpp

@@ -24,7 +24,10 @@ class state_base : public ::beman::task::detail::result_type<::beman::task::deta
 
     auto complete() -> std::coroutine_handle<> { return this->do_complete(); }
     auto get_stop_token() -> stop_token_type { return this->do_get_stop_token(); }
-    auto get_environment() -> Environment& { return this->do_get_environment(); }
+    auto get_environment() -> Environment& {
+        assert(this);
+        return this->do_get_environment();
+    }
     auto get_scheduler() -> scheduler_type { return this->do_get_scheduler(); }
     auto set_scheduler(scheduler_type other) -> scheduler_type { return this->do_set_scheduler(other); }
 

include/beman/task/detail/task.hpp

@@ -72,6 +72,7 @@ class task {
     }
     template <typename ParentPromise>
     auto as_awaitable(ParentPromise&) -> ::beman::task::detail::awaiter<Value, Env, promise_type, ParentPromise> {
+        assert(this->handle.get());
         return ::beman::task::detail::awaiter<Value, Env, promise_type, ParentPromise>(::std::move(this->handle));
     }
 
@@ -88,20 +89,13 @@ class task {
     struct domain {
         template <typename DS>
         auto transform_sender(DS&&, auto&&...) const noexcept {
-            std::cout << "task::domain::transform_sender\n";
             return dom_sender{};
         }
     };
     struct env {
-        auto query(const ::beman::execution::get_domain_t&) const noexcept -> domain {
-            std::cout << "task::env::get_domain\n";
-            return domain{};
-        }
+        auto query(const ::beman::execution::get_domain_t&) const noexcept -> domain { return domain{}; }
     };
-    auto xget_env() const noexcept {
-        std::cout << "task::get_env\n";
-        return env{};
-    }
+    auto xget_env() const noexcept { return env{}; }
 
   private:
     ::beman::task::detail::handle<promise_type> handle;