Improve handling of type-erased receiver environment.

Instead of having a dedicated parameter of type `env`, the receiver
environment is now exposed thorough `receiver`.
We don't need to store any extra data in the operation state, and the
receiver can be queried at any time for receiver environment properties.

Author: lucteo

include/exec/__detail/__system_context_default_impl.hpp

@@ -30,7 +30,6 @@ namespace exec::__system_context_default_impl {
   using system_context_replaceability::receiver;
   using system_context_replaceability::bulk_item_receiver;
   using system_context_replaceability::storage;
-  using system_context_replaceability::env;
   using system_context_replaceability::system_scheduler;
   using system_context_replaceability::__system_context_backend_factory;
 
@@ -45,20 +44,18 @@ namespace exec::__system_context_default_impl {
   - __recv::__r_ (receiver*) -- 8
   - __recv::__op_ (__operation*) -- 8
   - __operation::__inner_op_ (stdexec::connect_result_t<_Sender, __recv<_Sender>>) -- 56 (when connected with an empty receiver)
-  - __operation::__stop_token_ (inplace_stop_token) -- 8
   - __operation::__on_heap_ (bool) -- optimized away
   ---------------------
-  Total: 80; extra 24 bytes compared to internal operation state.
+  Total: 72; extra 16 bytes compared to internal operation state.
 
   extra for bulk:
   - __recv::__r_ (receiver*) -- 8
   - __recv::__op_ (__operation*) -- 8
   - __operation::__inner_op_ (stdexec::connect_result_t<_Sender, __recv<_Sender>>) -- 128 (when connected with an empty receiver & fun)
-  - __operation::__stop_token_ (inplace_stop_token) -- 8
   - __operation::__on_heap_ (bool) -- optimized away
   - __bulk_functor::__r_ (bulk_item_receiver*) - 8
   ---------------------
-  Total: 160; extra 32 bytes compared to internal operation state.
+  Total: 152; extra 24 bytes compared to internal operation state.
 
   Using libdispatch backend, the operation sizes are 48 (down from 80) and 128 (down from 160).
 
@@ -100,11 +97,9 @@ namespace exec::__system_context_default_impl {
     }
 
     decltype(auto) get_env() const noexcept {
-      return stdexec::prop{stdexec::get_stop_token, __get_stop_token()};
-    }
-
-    stdexec::inplace_stop_token __get_stop_token() const noexcept {
-      return __op_->__stop_token_;
+      auto __o = __r_->try_query<stdexec::inplace_stop_token>();
+      stdexec::inplace_stop_token __st = __o ? *__o : stdexec::inplace_stop_token{};
+      return stdexec::prop{stdexec::get_stop_token, __st};
     }
   };
 
@@ -127,22 +122,17 @@ namespace exec::__system_context_default_impl {
   struct __operation {
     /// The inner operation state, that results out of connecting the underlying sender with the receiver.
     stdexec::connect_result_t<_Sender, __recv<_Sender>> __inner_op_;
-    /// The stop token sent to the operation throught dynamic environment.
-    stdexec::inplace_stop_token __stop_token_;
     /// True if the operation is on the heap, false if it is in the preallocated space.
     bool __on_heap_;
 
     /// Try to construct the operation in the preallocated memory if it fits, otherwise allocate a new operation.
-    static __operation* __construct_maybe_alloc(
-      storage __storage,
-      receiver* __completion,
-      _Sender __sndr,
-      stdexec::inplace_stop_token __st) {
+    static __operation*
+      __construct_maybe_alloc(storage __storage, receiver* __completion, _Sender __sndr) {
       __storage = __ensure_alignment(__storage, alignof(__operation));
       if (__storage.__data == nullptr || __storage.__size < sizeof(__operation)) {
-        return new __operation(std::move(__sndr), __completion, __st, true);
+        return new __operation(std::move(__sndr), __completion, true);
       } else {
-        return new (__storage.__data) __operation(std::move(__sndr), __completion, __st, false);
+        return new (__storage.__data) __operation(std::move(__sndr), __completion, false);
       }
     }
 
@@ -161,13 +151,8 @@ namespace exec::__system_context_default_impl {
     }
 
    private:
-    __operation(
-      _Sender __sndr,
-      receiver* __completion,
-      stdexec::inplace_stop_token __st,
-      bool __on_heap)
+    __operation(_Sender __sndr, receiver* __completion, bool __on_heap)
       : __inner_op_(stdexec::connect(std::move(__sndr), __recv<_Sender>{__completion, this}))
-      , __stop_token_(std::move(__st))
       , __on_heap_(__on_heap) {
     }
   };
@@ -202,29 +187,25 @@ namespace exec::__system_context_default_impl {
       std::declval<__bulk_functor>()))>;
 
    public:
-    void schedule(storage __storage, receiver* __r, env __e) noexcept override {
+    void schedule(storage __storage, receiver* __r) noexcept override {
       try {
-        auto __o = __e.template try_query<stdexec::inplace_stop_token>();
-        auto __st = __o ? *__o : stdexec::inplace_stop_token{};
         auto __sndr = stdexec::schedule(__pool_scheduler_);
-        auto __os = __schedule_operation_t::__construct_maybe_alloc(
-          __storage, __r, std::move(__sndr), std::move(__st));
+        auto __os =
+          __schedule_operation_t::__construct_maybe_alloc(__storage, __r, std::move(__sndr));
         __os->start();
       } catch (std::exception& __e) {
         __r->set_error(std::current_exception());
       }
     }
 
     void
-      bulk_schedule(uint32_t __size, storage __storage, bulk_item_receiver* __r, env __e) noexcept
+      bulk_schedule(uint32_t __size, storage __storage, bulk_item_receiver* __r) noexcept
       override {
       try {
-        auto __o = __e.template try_query<stdexec::inplace_stop_token>();
-        auto __st = __o ? *__o : stdexec::inplace_stop_token{};
         auto __sndr =
           stdexec::bulk(stdexec::schedule(__pool_scheduler_), __size, __bulk_functor{__r});
-        auto __os = __bulk_schedule_operation_t::__construct_maybe_alloc(
-          __storage, __r, std::move(__sndr), std::move(__st));
+        auto __os =
+          __bulk_schedule_operation_t::__construct_maybe_alloc(__storage, __r, std::move(__sndr));
         __os->start();
       } catch (std::exception& __e) {
         __r->set_error(std::current_exception());

include/exec/__detail/__system_context_replaceability_api.hpp

@@ -78,12 +78,29 @@ namespace exec::system_context_replaceability {
   struct receiver {
     virtual ~receiver() = default;
 
+   protected:
+    virtual bool __query_env(__uuid, void*) noexcept = 0;
+
+   public:
     /// Called when the system scheduler completes successfully.
     virtual void set_value() noexcept = 0;
     /// Called when the system scheduler completes with an error.
     virtual void set_error(std::exception_ptr) noexcept = 0;
     /// Called when the system scheduler was stopped.
     virtual void set_stopped() noexcept = 0;
+
+    /// Query the receiver for a property of type `_P`.
+    template <typename _P>
+    std::optional<std::decay_t<_P>> try_query() noexcept {
+      if constexpr (__runtime_property<_P>) {
+        std::decay_t<_P> __p;
+        bool __success =
+          __query_env(__runtime_property_helper<std::decay_t<_P>>::__property_identifier, &__p);
+        return __success ? std::make_optional(std::move(__p)) : std::nullopt;
+      } else {
+        return std::nullopt;
+      }
+    }
   };
 
   /// Receiver for bulk sheduling operations.
@@ -99,82 +116,16 @@ namespace exec::system_context_replaceability {
     uint32_t __size;
   };
 
-  struct env {
-    /// Query the system context for a property of type `_P`.
-    template <typename _P>
-    std::optional<_P> try_query() noexcept {
-      if constexpr (__runtime_property<_P>) {
-        const void* __r = __query_(__data_, __runtime_property_helper<_P>::__property_identifier);
-        return __r ? std::make_optional(*static_cast<const _P*>(__r)) : std::nullopt;
-      } else {
-        return std::nullopt;
-      }
-    }
-
-    // IMPLEMENTATION DETAIL: used by the frontend with no proerties.
-    explicit env() noexcept
-      : __data_{nullptr}
-      , __query_{&__query_none} {
-    }
-
-    // IMPLEMENTATION DETAIL: used by the frontend with a single property.
-    template <__runtime_property _P>
-    explicit env(const _P& __p) noexcept
-      : __data_{std::addressof(__p)}
-      , __query_{&__query_single<_P>} {
-    }
-
-    // IMPLEMENTATION DETAIL: used by the frontend with multiple properties.
-    template <__runtime_property... _Ps>
-    explicit env(const std::tuple<_Ps...>& __properties) noexcept
-      : __data_{std::addressof(__properties)}
-      , __query_{&__query_multi<_Ps...>} {
-    }
-
-
-   private:
-    /// The source data containing all the properties.
-    const void* __data_;
-    /// Function called by the backend to query for a specific property.
-    const void* (*__query_)(const void*, __uuid) noexcept;
-
-    static const void* __query_none(const void* __data, __uuid __id) noexcept {
-      return nullptr;
-    }
-
-    template <__runtime_property _P>
-    static const void* __query_single(const void* __data, __uuid __id) noexcept {
-      if (__id == __runtime_property_helper<_P>::__property_identifier)
-        return __data;
-      return nullptr;
-    }
-
-    template <__runtime_property _P, typename _Tuple>
-    static uintptr_t __select_property_as_int(const _Tuple& __tuple, __uuid __id) noexcept {
-      if (__id == __runtime_property_helper<_P>::__property_identifier)
-        return reinterpret_cast<uintptr_t>(std::addressof(std::get<_P>(__tuple)));
-      return 0;
-    }
-
-    template <__runtime_property... _Ps>
-    static const void* __query_multi(const void* __data, __uuid __id) noexcept {
-      const std::tuple<_Ps...>& __properties = *static_cast<const std::tuple<_Ps...>*>(__data);
-      uintptr_t __result = (... + __select_property_as_int<_Ps>(__properties, __id));
-      return reinterpret_cast<const void*>(__result);
-    }
-  };
-
   /// Interface for the system scheduler
   struct system_scheduler {
     static constexpr __uuid __interface_identifier{0x5ee9202498c4bd4f, 0xa1df2508ffcd9d7e};
 
     virtual ~system_scheduler() = default;
 
-    /// Schedule work on system scheduler, calling `__r` when done and using `__s` for preallocated memory, using `__e` for environment.
-    virtual void schedule(storage __s, receiver* __r, env __e) noexcept = 0;
-    /// Schedule bulk work of size `__n` on system scheduler, calling `__r` for each item and then when done, and using `__s` for preallocated memory, using `__e` for environment.
-    virtual void
-      bulk_schedule(uint32_t __n, storage __s, bulk_item_receiver* __r, env __e) noexcept = 0;
+    /// Schedule work on system scheduler, calling `__r` when done and using `__s` for preallocated memory.
+    virtual void schedule(storage __s, receiver* __r) noexcept = 0;
+    /// Schedule bulk work of size `__n` on system scheduler, calling `__r` for each item and then when done, and using `__s` for preallocated memory.
+    virtual void bulk_schedule(uint32_t __n, storage __s, bulk_item_receiver* __r) noexcept = 0;
   };
 
 } // namespace exec::system_context_replaceability

include/exec/system_context.hpp

@@ -44,6 +44,20 @@ namespace exec {
         : __rcvr_{std::forward<_Rcvr>(__rcvr)} {
       }
 
+      bool __query_env(__uuid __id, void* __dest) noexcept override {
+        using system_context_replaceability::__runtime_property_helper;
+        using __StopToken = decltype(stdexec::get_stop_token(stdexec::get_env(__rcvr_)));
+        if constexpr (std::is_same_v<stdexec::inplace_stop_token, __StopToken>) {
+          if (
+            __id == __runtime_property_helper<stdexec::inplace_stop_token>::__property_identifier) {
+            *static_cast<stdexec::inplace_stop_token*>(__dest) =
+              stdexec::get_stop_token(stdexec::get_env(__rcvr_));
+            return true;
+          }
+        }
+        return false;
+      }
+
       void set_value() noexcept override {
         stdexec::set_value(std::forward<_Rcvr>(__rcvr_));
       }
@@ -174,14 +188,10 @@ namespace exec {
           stdexec::set_stopped(__rcvr_);
           return;
         }
-        system_context_replaceability::env __e{};
-        if constexpr (std::is_same_v<stdexec::inplace_stop_token, std::decay_t<decltype(st)>>) {
-          __e = system_context_replaceability::env{st};
-        }
         auto& __scheduler_impl = __preallocated_.__as<__system_scheduler_ptr>();
         auto __impl = std::move(__scheduler_impl);
         std::destroy_at(&__scheduler_impl);
-        __impl->schedule(__preallocated_.__as_storage(), &__rcvr_, __e);
+        __impl->schedule(__preallocated_.__as_storage(), &__rcvr_);
       }
 
       /// Object that receives completion from the work described by the sender.
@@ -304,6 +314,21 @@ namespace exec {
           std::tuple<stdexec::__decay_t<_As>...>{std::move(__as)...};
       }
 
+      bool __query_env(__uuid __id, void* __dest) noexcept override {
+        auto __state = reinterpret_cast<_BulkState*>(this);
+        using system_context_replaceability::__runtime_property_helper;
+        using __StopToken = decltype(stdexec::get_stop_token(stdexec::get_env(__state->__rcvr_)));
+        if constexpr (std::is_same_v<stdexec::inplace_stop_token, __StopToken>) {
+          if (
+            __id == __runtime_property_helper<stdexec::inplace_stop_token>::__property_identifier) {
+            *static_cast<stdexec::inplace_stop_token*>(__dest) =
+              stdexec::get_stop_token(stdexec::get_env(__state->__rcvr_));
+            return true;
+          }
+        }
+        return false;
+      }
+
       /// Calls `set_value()` on the final receiver of the bulk operation, using the values from the previous sender.
       void set_value() noexcept override {
         auto __state = reinterpret_cast<_BulkState*>(this);
@@ -385,10 +410,6 @@ namespace exec {
           stdexec::set_stopped(__state_.__rcvr_);
           return;
         }
-        system_context_replaceability::env __e{};
-        if constexpr (std::is_same_v<stdexec::inplace_stop_token, std::decay_t<decltype(st)>>) {
-          __e = system_context_replaceability::env{st};
-        }
 
         // Store the input data in the shared state.
         using __typed_forward_args_receiver_t =
@@ -404,7 +425,7 @@ namespace exec {
 
         // Schedule the bulk work on the system scheduler.
         // This will invoke `start` on our receiver multiple times, and then a completion signal (e.g., `set_value`).
-        __scheduler->bulk_schedule(__size, __storage, __r, __e);
+        __scheduler->bulk_schedule(__size, __storage, __r);
       }
 
       /// Invoked when the previous sender completes with "stopped" to stop the entire work.