From b0859bc640dbd1cc8b92dbe40c14a51bd620b6ce Mon Sep 17 00:00:00 2001
From: Maikel Nadolski <maikel.nadolski@gmail.com>
Date: Sun, 24 Sep 2023 22:14:35 +0200
Subject: [PATCH] Implement parallelism and cardinality

---
 include/exec/sequence/empty_sequence.hpp |  7 ++
 include/exec/sequence/iterate.hpp        | 11 +++-
 include/exec/sequence_senders.hpp        | 84 +++++++++++++++++++++++-
 test/exec/sequence/test_iterate.cpp      | 11 ++++
 4 files changed, 110 insertions(+), 3 deletions(-)

diff --git a/include/exec/sequence/empty_sequence.hpp b/include/exec/sequence/empty_sequence.hpp
index 0be237b62..05f8273c1 100644
--- a/include/exec/sequence/empty_sequence.hpp
+++ b/include/exec/sequence/empty_sequence.hpp
@@ -50,6 +50,13 @@ namespace exec {
           return stdexec::__t<__operation<stdexec::__id<__decay_t<_Rcvr>>>>{
             static_cast<_Rcvr&&>(__rcvr)};
         }
+
+        template <__decays_to<__t> _Self>
+        friend auto tag_invoke(get_env_t, _Self&&) noexcept {
+          return make_env(
+            with(parallelism, lock_step),
+            with(cardinality, std::integral_constant<std::size_t, 0>{}));
+        }
       };
     };
 
diff --git a/include/exec/sequence/iterate.hpp b/include/exec/sequence/iterate.hpp
index 96ef4550f..6107a3a10 100644
--- a/include/exec/sequence/iterate.hpp
+++ b/include/exec/sequence/iterate.hpp
@@ -205,8 +205,15 @@ namespace exec {
         return {};
       }
 
-      static empty_env get_env(__ignore) noexcept {
-        return {};
+      static make_env_t<with_t<parallelism_t, lock_step_t>> get_env(__ignore) noexcept {
+        return make_env(with(parallelism, lock_step));
+      }
+
+      template <sender_expr_for<iterate_t> _SeqExpr>
+        requires std::ranges::sized_range<stdexec::__data_of<_SeqExpr>>
+      static auto get_env(const _SeqExpr& __seq) noexcept {
+        auto&& __rng = apply_sender(__seq, stdexec::__detail::__get_data{});
+        return make_env(with(parallelism, lock_step), with(cardinality, std::ranges::size(__rng)));
       }
     };
   }
diff --git a/include/exec/sequence_senders.hpp b/include/exec/sequence_senders.hpp
index 5dbc92cc8..f5792282b 100644
--- a/include/exec/sequence_senders.hpp
+++ b/include/exec/sequence_senders.hpp
@@ -18,6 +18,8 @@
 
 #include "../stdexec/execution.hpp"
 
+#include "./env.hpp"
+
 namespace exec {
   struct sequence_tag { };
 
@@ -193,7 +195,7 @@ namespace exec {
   using item_types_of_t =
     decltype(get_item_types(stdexec::__declval<_Sender>(), stdexec::__declval<_Env>()));
 
-  template <class _Sender, class _Env>
+  template <class _Sender, class _Env = stdexec::no_env>
   concept sequence_sender =           //
     stdexec::sender<_Sender, _Env> && //
     enable_sequence_sender<stdexec::__decay_t<_Sender>>;
@@ -454,4 +456,84 @@ namespace exec {
       }
     }
   }
+
+  namespace __sequence_queries {
+    using namespace stdexec;
+
+    inline struct unbounded_t {
+    } unbounded;
+
+    struct cardinality_t {
+      template <class Env>
+        requires tag_invocable<cardinality_t, const Env&>
+      constexpr tag_invoke_result_t<cardinality_t, const Env&> operator()(const Env& env) const
+        noexcept(nothrow_tag_invocable<cardinality_t, const Env&>) {
+        return tag_invoke(*this, env);
+      }
+
+      friend constexpr bool tag_invoke(forwarding_query_t, cardinality_t) noexcept {
+        return true;
+      }
+    };
+
+    inline struct many_sender_t {
+    } many_sender;
+
+    inline struct lock_step_t {
+    } lock_step;
+
+    struct parallelism_t {
+      template <class Env>
+        requires tag_invocable<parallelism_t, const Env&>
+      constexpr tag_invoke_result_t<parallelism_t, const Env&> operator()(const Env& env) const
+        noexcept(nothrow_tag_invocable<parallelism_t, const Env&>) {
+        return tag_invoke(*this, env);
+      }
+
+      friend constexpr bool tag_invoke(forwarding_query_t, parallelism_t) noexcept {
+        return true;
+      }
+    };
+  }
+
+  using __sequence_queries::cardinality_t;
+  inline constexpr cardinality_t cardinality;
+
+  using __sequence_queries::parallelism_t;
+  inline constexpr parallelism_t parallelism;
+
+  using __sequence_queries::unbounded_t;
+  using __sequence_queries::unbounded;
+
+  using __sequence_queries::lock_step_t;
+  using __sequence_queries::lock_step;
+
+  using __sequence_queries::many_sender_t;
+  using __sequence_queries::many_sender;
+
+  namespace __get_sequence_env {
+    using namespace stdexec;
+
+    struct get_sequence_env_t {
+      template <sequence_sender Sequence>
+      constexpr env_of_t<Sequence> operator()(const Sequence& seq) const noexcept {
+        return stdexec::get_env(seq);
+      }
+
+      template <sender Sequence>
+        requires(!sequence_sender<Sequence>)
+      constexpr auto operator()(const Sequence& seq) const noexcept {
+        return make_env(
+          stdexec::get_env(seq),
+          with(cardinality, std::integral_constant<size_t, 1>{}),
+          with(parallelism, lock_step));
+      }
+    };
+  }
+
+  using __get_sequence_env::get_sequence_env_t;
+  inline constexpr get_sequence_env_t get_sequence_env;
+
+  template <class EnvProvider>
+  using sequence_env_of_t = stdexec::__call_result_t<get_sequence_env_t, EnvProvider>;
 }
\ No newline at end of file
diff --git a/test/exec/sequence/test_iterate.cpp b/test/exec/sequence/test_iterate.cpp
index bc867e4d2..222ada231 100644
--- a/test/exec/sequence/test_iterate.cpp
+++ b/test/exec/sequence/test_iterate.cpp
@@ -129,4 +129,15 @@ TEST_CASE("iterate - sum up an array with custom domain", "[sequence_senders][it
   CHECK(sum == (42 + 43 + 44 + 1));
 }
 
+TEST_CASE("iterate - is lock step") {
+  std::array<int, 3> array{42, 43, 44};
+  auto iterate = exec::iterate(std::views::all(array));
+  STATIC_REQUIRE(exec::sequence_sender_in<decltype(iterate), stdexec::empty_env>);
+  STATIC_REQUIRE(stdexec::sender_expr_for<decltype(iterate), exec::iterate_t>);
+  using parallelism_t = decltype(exec::parallelism(stdexec::get_env(iterate)));
+  STATIC_REQUIRE(std::same_as<parallelism_t, exec::lock_step_t>);
+  auto size = exec::cardinality(stdexec::get_env(iterate));
+  CHECK(size == 3);
+}
+
 #endif // STDEXEC_HAS_STD_RANGES()