Merge pull request #2248 from Cra3z/fix-starts_on

Fix starts_on's sender-attributes

Author: ericniebler

include/stdexec/__detail/__starts_on.hpp

@@ -65,32 +65,9 @@ namespace STDEXEC
         , __attr_(__attr)
       {}
 
-      // Query for completion scheduler to use for algorithm dispatching.
-      // NOT TO SPEC
-      template <class... _Env>
-      STDEXEC_ATTRIBUTE(nodiscard, always_inline, host, device)
-      constexpr auto query(get_completion_domain_t<>, _Env&&...) const noexcept
-        -> __completion_domain_of_t<set_value_t, _Scheduler, _Env...>
-      {
-        return {};
-      }
-
-      // Query for completion scheduler
-      template <class _SetTag, class... _Env>
-        requires __completes_where_it_starts<_SetTag, env_of_t<_Child>, __env2_t<_Env>...>
-      STDEXEC_ATTRIBUTE(nodiscard, always_inline, host, device)
-      constexpr auto query(get_completion_scheduler_t<_SetTag>, _Env&&...) const noexcept
-        -> _Scheduler
-      {
-        // If the child completes where it starts, then starts_on(Sch,Child) completes on
-        // scheduler Sch.
-        return __sched_;
-      }
-
       // Query for completion scheduler - delegates to child's env with augmented
       // environment
       template <class _SetTag, class... _Env>
-        requires(!__completes_where_it_starts<_SetTag, env_of_t<_Child>, __env2_t<_Env>...>)
       STDEXEC_ATTRIBUTE(nodiscard, always_inline, host, device)
       constexpr auto query(get_completion_scheduler_t<_SetTag> __query,
                            _Env&&... __env) const noexcept

test/stdexec/algos/adaptors/test_starts_on.cpp

@@ -319,6 +319,58 @@ namespace
     ex::sender auto                snd = ex::starts_on(is, ex::just()) | ex::then([] {});
     ex::sync_wait(std::move(snd));
   }
+
+  // A sender that completes asynchronously on the same execution context that started it
+  struct affine_sender
+  {
+    using sender_concept        = ex::sender_tag;
+    using completion_signatures = ex::completion_signatures<ex::set_value_t()>;
+
+    struct attrs
+    {
+      template <class Tag, class... Env>
+      static auto query(ex::__get_completion_behavior_t<Tag>, Env&&...) noexcept
+      {
+        return ex::__completion_behavior::__asynchronous_affine;
+      }
+    };
+
+    template <class Receiver>
+    static auto connect(Receiver rcvr)
+    {
+      return ex::connect(ex::just(), static_cast<Receiver&&>(rcvr));
+    }
+
+    [[nodiscard]]
+    static auto get_env() noexcept -> attrs
+    {
+      return {};
+    }
+  };
+
+  TEST_CASE("starts_on completion domain agrees with its completion scheduler when the child "
+            "completes where it starts",
+            "[adaptors][starts_on]")
+  {
+    exec::static_thread_pool pool{1};
+    auto                     sched = pool.get_scheduler();
+    using sched_t                  = decltype(sched);
+    using sched_domain_t =
+      ex::__call_result_t<ex::get_completion_domain_t<ex::set_value_t>, sched_t, ex::env<>>;
+
+    auto snd      = ex::starts_on(sched, affine_sender{});
+    using attrs_t = ex::env_of_t<decltype(snd)>;
+
+    using cmpl_sched_t =
+      ex::__call_result_t<ex::get_completion_scheduler_t<ex::set_value_t>, attrs_t, ex::env<>>;
+    using cmpl_domain_t =
+      ex::__call_result_t<ex::get_completion_domain_t<ex::set_value_t>, attrs_t, ex::env<>>;
+    STATIC_REQUIRE(std::same_as<cmpl_sched_t, sched_t>);
+    STATIC_REQUIRE(std::same_as<cmpl_domain_t, sched_domain_t>);
+    CHECK(ex::get_completion_scheduler<ex::set_value_t>(ex::get_env(snd), ex::env<>{}) == sched);
+
+    ex::sync_wait(std::move(snd));
+  }
 }  // namespace
 
 STDEXEC_PRAGMA_POP()

test/stdexec/types/test_task.cpp

@@ -425,6 +425,31 @@ namespace
     CHECK(i == 42);
   }
 
+  auto noop_task() -> ex::task<>
+  {
+    co_return;
+  }
+
+  auto test_task_awaits_starts_on_task() -> ex::task<>
+  {
+    auto sched = co_await ex::read_env(ex::get_start_scheduler);
+    co_await ex::starts_on(sched, noop_task());
+  }
+
+  TEST_CASE("test starts_on(task_scheduler, task) advertises task_scheduler_domain and can be "
+            "sync_wait'ed",
+            "[types][task]")
+  {
+    exec::single_thread_context ctx;
+    ex::task_scheduler          sched{ctx.get_scheduler()};
+    auto                        snd = ex::starts_on(sched, test_task_awaits_starts_on_task());
+    using attrs_t                   = ex::env_of_t<decltype(snd)>;
+    using domain_t =
+      ex::__call_result_t<ex::get_completion_domain_t<ex::set_value_t>, attrs_t, ex::env<>>;
+    STATIC_REQUIRE(std::same_as<domain_t, ex::task_scheduler_domain>);
+    CHECK(ex::sync_wait(std::move(snd)).has_value());
+  }
+
   // Test affinity with a run_loop scheduler, which is infallible but not inline:
   struct test_env2
   {