From 097b73250b54c420899ea3845609a54532d31cc1 Mon Sep 17 00:00:00 2001
From: Eric Niebler <eniebler@nvidia.com>
Date: Thu, 3 Sep 2026 09:30:47 -0700
Subject: [PATCH] use start sched/domain as completion sched/domain for senders
 that are affine, not just inline

---
 include/stdexec/__detail/__domain.hpp     | 37 +++++++++++++----------
 include/stdexec/__detail/__schedulers.hpp |  6 ++--
 test/stdexec/types/test_task.cpp          | 14 +++++++++
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/include/stdexec/__detail/__domain.hpp b/include/stdexec/__detail/__domain.hpp
index e2a441675..be33b55bf 100644
--- a/include/stdexec/__detail/__domain.hpp
+++ b/include/stdexec/__detail/__domain.hpp
@@ -243,8 +243,8 @@ namespace STDEXEC
   STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
   __hide_query(_Env &&, _Queries...) -> __hide_query<_Env, _Queries...>;
 
-  //! @brief A wrapper around an environment that hides the get_start_scheduler and get_domain
-  //! queries.
+  //! @brief A wrapper around an environment that hides the get_start_scheduler and
+  //! get_domain queries.
   template <class _Env>
   struct __hide_scheduler : __hide_query<_Env, get_start_scheduler_t>
   {
@@ -304,7 +304,8 @@ namespace STDEXEC
     static consteval auto __check_domain(_Domain) noexcept -> _Domain
     {
       // Sanity check: if a completion scheduler can be determined from the attributes
-      // (not the environment), then its domain must match the domain returned by the attributes.
+      // (not the environment), then its domain must match the domain returned by the
+      // attributes.
       if constexpr (!__same_as<_Tag, void>)
       {
         if constexpr (__callable<get_completion_scheduler_t<_Tag>, _Attrs const &, _Env const &...>)
@@ -330,8 +331,9 @@ namespace STDEXEC
       {
         using __domain_t = __call_result_t<__read_query_t, _Attrs const &, _Env const &...>;
         return __check_domain<_Attrs, _Env...>(__domain_t{});
-        // Otherwise, if _Tag is void, fall back to querying for the set_value_t completion domain:
       }
+      // Otherwise, if _Tag is void, fall back to querying for the set_value_t completion
+      // domain:
       else if constexpr (__same_as<_Tag, void>)
       {
         if constexpr (__callable<get_completion_domain_t<set_value_t>,
@@ -347,8 +349,8 @@ namespace STDEXEC
           return void();
         }
       }
-      // Otherwise, if __attrs has a completion scheduler, we can ask that scheduler for its
-      // completion domain.
+      // Otherwise, if __attrs has a completion scheduler, we can ask that scheduler for
+      // its completion domain.
       else if constexpr (__callable<get_completion_scheduler_t<_Tag>,
                                     _Attrs const &,
                                     _Env const &...>)
@@ -362,29 +364,32 @@ namespace STDEXEC
           using __domain_t = __call_result_t<__read_query_t, __sch_t, _Env const &...>;
           return __domain_t{};
         }
-        // Otherwise, if the scheduler's sender indicates that it completes inline, we can ask
-        // the environment for its domain.
-        else if constexpr (__completes_inline<_Tag,
-                                              env_of_t<__call_result_t<schedule_t, __sch_t>>,
-                                              _Env...>
+        // Otherwise, if the scheduler's sender indicates that it completes where it
+        // starts, we can ask the environment for its domain.
+        else if constexpr (__completes_where_it_starts<
+                             _Tag,
+                             env_of_t<__call_result_t<schedule_t, __sch_t>>,
+                             _Env...>
                            && __callable<get_domain_t, _Env const &...>)
         {
           return __call_result_t<get_domain_t, _Env const &...>{};
         }
-        // Otherwise, if we are asking "late" (with an environment), return the default_domain
+        // Otherwise, if we are asking "late" (with an environment), return the
+        // default_domain
         else if constexpr (sizeof...(_Env) != 0)
         {
           return default_domain{};
         }
       }
-      // Otherwise, if the attributes indicates that the sender completes inline, we can ask
-      // the environment for its domain.
-      else if constexpr (__completes_inline<_Tag, _Attrs, _Env...>
+      // Otherwise, if the attributes indicates that the sender completes where it starts,
+      // we can ask the environment for its domain.
+      else if constexpr (__completes_where_it_starts<_Tag, _Attrs, _Env...>
                          && __callable<get_domain_t, _Env const &...>)
       {
         return __call_result_t<get_domain_t, _Env const &...>{};
       }
-      // Otherwise, if we are asking "late" (with an environment), return the default_domain
+      // Otherwise, if we are asking "late" (with an environment), return the
+      // default_domain
       else if constexpr (sizeof...(_Env) != 0)
       {
         return default_domain{};
diff --git a/include/stdexec/__detail/__schedulers.hpp b/include/stdexec/__detail/__schedulers.hpp
index 7c73c6be6..2fe9f4c9e 100644
--- a/include/stdexec/__detail/__schedulers.hpp
+++ b/include/stdexec/__detail/__schedulers.hpp
@@ -427,7 +427,7 @@ namespace STDEXEC
       // Otherwise, if __attrs indicates that its sender completes inline, then we can ask
       // the environment for the current scheduler and return that (after checking the
       // scheduler for _its_ completion scheduler).
-      else if constexpr (__completes_inline<_Tag, _Attrs, _Env...>
+      else if constexpr (__completes_where_it_starts<_Tag, _Attrs, _Env...>
                          && (__callable<get_start_scheduler_t, _Env const &> || ...))
       {
         using __result_t = __call_result_t<__recurse_query_t,
@@ -461,7 +461,7 @@ namespace STDEXEC
       // Otherwise, if __attrs indicates that its sender completes inline, then we can ask
       // the environment for the current scheduler and return that (after checking the
       // scheduler for _its_ completion scheduler).
-      else if constexpr (__completes_inline<_Tag, _Attrs, _Env...>
+      else if constexpr (__completes_where_it_starts<_Tag, _Attrs, _Env...>
                          && __callable<get_start_scheduler_t, _Env const &...>)
       {
         return __check_domain<_Attrs, _Env...>(
@@ -604,7 +604,7 @@ namespace STDEXEC
   template <class _Sch, class _Env>
   constexpr auto __mk_sch_env([[maybe_unused]] _Sch __sch, _Env const &__env) noexcept
   {
-    if constexpr (__completes_inline<set_value_t, env_of_t<schedule_result_t<_Sch>>, _Env>
+    if constexpr (__completes_where_it_starts<set_value_t, env_of_t<schedule_result_t<_Sch>>, _Env>
                   && __callable<get_start_scheduler_t, _Env const &>)
     {
       auto __sch2 = get_completion_scheduler<set_value_t>(get_start_scheduler(__env),
diff --git a/test/stdexec/types/test_task.cpp b/test/stdexec/types/test_task.cpp
index 22e5c4fda..71a9b2558 100644
--- a/test/stdexec/types/test_task.cpp
+++ b/test/stdexec/types/test_task.cpp
@@ -654,6 +654,20 @@ namespace
     CHECK(!res.has_value());
   }
 
+  TEST_CASE("test completion domain of task", "[types][task]")
+  {
+    // task is scheduler affine but not inline. Regardless, its completion domain is the
+    // same as the start scheduler's completion domain.
+    using attrs_t = ex::env_of_t<ex::task<int>>;
+    using env_t   = ex::prop<ex::get_start_scheduler_t, ex::parallel_scheduler>;
+    using sched_t =
+      std::invoke_result_t<ex::get_completion_scheduler_t<ex::set_value_t>, attrs_t, env_t>;
+    using domain_t =
+      std::invoke_result_t<ex::get_completion_domain_t<ex::set_value_t>, attrs_t, env_t>;
+    STATIC_REQUIRE(std::same_as<sched_t, ex::parallel_scheduler>);
+    STATIC_REQUIRE(std::same_as<domain_t, ex::__parallel_scheduler_domain>);
+  }
+
   TEST_CASE("repro for NVIDIA/stdexec#2041", "[types][task]")
   {
     auto task = []() -> ex::task<void>