Add stdexec::simple_counting_scope

This diff adds `stdexec::simple_counting_scope` and one compeletely
trivial test case. Still a work in progress.

Author: ispeters

include/stdexec/__detail/__counting_scopes.hpp

@@ -0,0 +1,395 @@
+/*
+ * Copyright (c) 2025 Ian Petersen
+ * Copyright (c) 2025 NVIDIA Corporation
+ *
+ * Licensed under the Apache License Version 2.0 with LLVM Exceptions
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *   https://llvm.org/LICENSE.txt
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "__execution_fwd.hpp"
+
+#include "__concepts.hpp"
+#include "__env.hpp"
+#include "__receivers.hpp"
+#include "__schedulers.hpp"
+#include "__senders_core.hpp"
+#include "__sender_introspection.hpp"
+#include "__type_traits.hpp"
+
+#include <atomic>
+#include <cstddef>
+#include <exception>
+#include <limits>
+#include <memory>
+#include <type_traits>
+#include <utility>
+
+namespace stdexec {
+  class simple_counting_scope;
+
+  namespace __counting_scopes {
+    struct __state_base {
+      __state_base* __next_{};
+
+      __state_base() = default;
+      __state_base(__state_base&&) = delete;
+
+      virtual void __complete() noexcept = 0;
+
+     protected:
+      ~__state_base() = default;
+    };
+
+    struct __scope_join_t { };
+
+    struct __scope_join_impl : __sexpr_defaults {
+      template <class _Scope, class _Rcvr>
+      struct __state final : __state_base {
+        struct __rcvr_t {
+          using receiver_concept = receiver_t;
+
+          _Rcvr& __rcvr_;
+
+          void set_value() && noexcept {
+            stdexec::set_value(std::move(__rcvr_));
+          }
+
+          template <class E>
+          void set_error(E&& e) && noexcept {
+            stdexec::set_error(std::move(__rcvr_), static_cast<E&&>(e));
+          };
+
+          void set_stopped() && noexcept {
+            stdexec::set_stopped(std::move(__rcvr_));
+          }
+
+          decltype(auto) get_env() const noexcept {
+            return stdexec::get_env(__rcvr_);
+          }
+        };
+
+        using __sched_sender = decltype(schedule(get_scheduler(get_env(__declval<_Rcvr&>()))));
+
+        using __op_t = connect_result_t<__sched_sender, __rcvr_t>;
+
+        _Scope* __scope_;
+        _Rcvr& __receiver_;
+        __op_t __op_;
+
+        __state(_Scope* __scope, _Rcvr& __rcvr)
+          noexcept(__nothrow_callable<connect_t, __sched_sender, __rcvr_t>)
+          : __scope_(__scope)
+          , __receiver_(__rcvr)
+          , __op_(connect(schedule(get_scheduler(get_env(__rcvr)))), __rcvr_t(__rcvr)) {
+        }
+
+        void __complete() noexcept override {
+          start(__op_);
+        }
+
+        void __complete_inline() noexcept {
+          set_value(std::move(__receiver_));
+        }
+      };
+
+      static constexpr auto get_state =
+        []<class _Sndr, class _Rcvr>(_Sndr&& __sender, _Rcvr& __receiver) noexcept(
+          __nothrow_constructible_from<
+            __state<_Rcvr, __data_of<std::remove_cvref_t<_Sndr>>>,
+            __data_of<std::remove_cvref_t<_Sndr>>,
+            _Rcvr&
+          >) {
+          auto [_, self] = __sender;
+          return __state(self, __receiver);
+        };
+
+      static constexpr auto start = [](auto& __s, auto&) noexcept {
+        if (__s.__scope_->__start_join_sender(__s)) {
+          __s.__complete_inline();
+        }
+      };
+    };
+
+    template <class _Scope>
+    struct __association_t {
+      constexpr __association_t() = default;
+
+      constexpr __association_t(__association_t&& __other) noexcept
+        : __scope_(std::exchange(__other.__scope_, nullptr)) {
+      }
+
+      ~__association_t() {
+        if (__scope_ != nullptr) {
+          __scope_->__disassociate();
+        }
+      }
+
+      __association_t& operator=(__association_t __rhs) noexcept {
+        std::swap(__scope_, __rhs.__scope_);
+        return *this;
+      }
+
+      constexpr explicit operator bool() const noexcept {
+        return __scope_ != nullptr;
+      }
+
+      __association_t try_associate() const noexcept {
+        if (__scope_) {
+          return __scope_->__try_associate();
+        } else {
+          return __association_t();
+        }
+      }
+
+     private:
+      friend simple_counting_scope;
+
+      _Scope* __scope_{};
+
+      constexpr __association_t(_Scope& __scope) noexcept
+        : __scope_(std::addressof(__scope)) {
+      }
+    };
+  } // namespace __counting_scopes
+
+  template <>
+  struct __sexpr_impl<__counting_scopes::__scope_join_t> : __counting_scopes::__scope_join_impl { };
+
+  class simple_counting_scope {
+   public:
+    using __assoc_t = __counting_scopes::__association_t<simple_counting_scope>;
+
+    struct token {
+      template <sender _Sender>
+      _Sender&& wrap(_Sender&& __snd) const noexcept {
+        return static_cast<_Sender&&>(__snd);
+      }
+
+      __assoc_t try_associate() const noexcept {
+        return __scope_->__try_associate();
+      }
+
+     private:
+      friend class simple_counting_scope;
+
+      simple_counting_scope* __scope_;
+
+      explicit token(simple_counting_scope* __scope) noexcept
+        : __scope_(__scope) {
+      }
+    };
+
+    static constexpr std::size_t max_associations = std::numeric_limits<std::size_t>::max() >> 3;
+
+    simple_counting_scope() = default;
+
+    simple_counting_scope(simple_counting_scope&&) = delete;
+
+    ~simple_counting_scope() {
+      auto state = __state_.load(std::memory_order_relaxed);
+      if (__is_join_needed(state) || __count(state) != 0ul) {
+        std::terminate();
+      }
+    }
+
+    token get_token() noexcept {
+      return token{this};
+    }
+
+    void close() noexcept {
+      __state_.fetch_or(__closed, std::memory_order_relaxed);
+    }
+
+    sender auto join() noexcept {
+      return __make_sexpr<__counting_scopes::__scope_join_t>(this);
+    }
+
+   private:
+    friend __assoc_t;
+
+    static constexpr std::size_t __closed{1ul};
+    static constexpr std::size_t __join_needed{2ul};
+    static constexpr std::size_t __join_running{4ul};
+
+    std::atomic<std::size_t> __state_{0ul};
+    std::atomic<void*> __waitingJoinOps_{nullptr};
+
+    __assoc_t __try_associate() noexcept {
+      auto state = __state_.load(std::memory_order_relaxed);
+
+      do {
+        if (__is_closed(state) || __count(state) == max_associations) {
+          // the scope is closed or full so deny a new association
+          return __assoc_t();
+        }
+
+        // increment the count and ensure the join-needed bit is set
+        const auto newState = __make_state(__count(state) + 1ul, __bits(state) | __join_needed);
+
+        assert(__count(newState) <= max_associations);
+
+        if (__state_.compare_exchange_weak(
+              state,
+              newState,
+              // this is effectively a ref-count increment so
+              // there's no need to synchronize
+              std::memory_order_relaxed)) {
+          return __assoc_t{*this};
+        }
+      } while (true);
+    }
+
+    void __disassociate() noexcept {
+      auto state = __state_.load(std::memory_order_relaxed);
+
+      assert(__count(state) > 0ul);
+
+      do {
+        const auto newCount = __count(state) - 1ul;
+        const auto newBits = (newCount == 0ul && __is_joining(state) ? __closed : __bits(state));
+
+        if (__state_.compare_exchange_weak(
+              state,
+              __make_state(newCount, newBits),
+              std::memory_order_acq_rel,
+              std::memory_order_relaxed)) {
+          // successfully updated; now decide whether we need to complete
+          // outstanding join-senders or just bail out
+          if (newCount == 0ul && newBits == __closed) {
+            // launch the outstanding join-senders after the loop
+            break;
+          } else {
+            // just bail out--either we're not the last op, or the scope
+            // isn't closed yet
+            return;
+          }
+        }
+      } while (true);
+
+      // mark the linked list as having been consumed
+      auto* joinOpsToComplete = static_cast<__counting_scopes::__state_base*>(
+        __waitingJoinOps_.exchange(this, std::memory_order_acq_rel));
+
+      for (; joinOpsToComplete != nullptr; joinOpsToComplete = joinOpsToComplete->__next_) {
+        joinOpsToComplete->__complete();
+      }
+    }
+
+    constexpr static bool __is_unused(std::size_t __state) noexcept {
+      return (__state & __join_needed) == 0ul;
+    }
+
+    constexpr static bool __is_open(std::size_t __state) noexcept {
+      return (__state & __closed) == 0ul;
+    }
+
+    constexpr static bool __is_closed(std::size_t __state) noexcept {
+      return !__is_open(__state);
+    }
+
+    constexpr static bool __is_joined(std::size_t __state) noexcept {
+      return __is_closed(__state) && !__is_unused(__state);
+    }
+
+    constexpr static bool __is_joining(std::size_t __state) noexcept {
+      return (__state & __join_running) != 0ul;
+    }
+
+    constexpr static bool __is_join_needed(std::size_t __state) noexcept {
+      return !__is_unused(__state);
+    }
+
+    constexpr static std::size_t __count(std::size_t __state) noexcept {
+      return __state >> 3;
+    }
+
+    constexpr static std::size_t __bits(std::size_t __state) noexcept {
+      return __state & 7ul;
+    }
+
+    constexpr static std::size_t
+      __make_state(std::size_t __opCount, std::size_t __lowBits) noexcept {
+      // no high bits set
+      assert(__count(__opCount << 3) == __opCount);
+
+      // no high bits set
+      assert(__bits(__lowBits) == __lowBits);
+
+      return (__opCount << 3) | __bits(__lowBits);
+    }
+
+    bool __start_join_sender(__counting_scopes::__state_base& __joinOp) noexcept {
+      auto state = __state_.load(std::memory_order_relaxed);
+
+      do {
+        // [exec.simple.counting.mem] para (9.1)
+        // unused, unused-and-closed, or joined -> joined
+        if (__is_unused(state) || __is_closed(state)) {
+          assert(__count(state) == 0ul);
+
+          const auto newState = __make_state(__count(state), __closed);
+
+          assert(__is_joined(newState));
+
+          // try to make it joined
+          if (__state_.compare_exchange_weak(
+                state, newState, std::memory_order_acq_rel, std::memory_order_relaxed)) {
+            return true;
+          }
+        }
+        // [exec.simple.counting.mem] para (9.2)
+        // open or open-and-joining -> open-and-joining
+        // [exec.simple.counting.mem] para (9.3)
+        // closed or closed-and-joining -> closed-and-joining
+        else {
+          assert(__is_join_needed(state));
+
+          // try to make it {open|closed}-and-joining
+          const auto newState = state | __join_running;
+
+          assert(__is_joining(newState));
+
+          // TODO: does this need to do any synchronizing or is relaxed OK?
+          if (__state_.compare_exchange_weak(state, newState, std::memory_order_relaxed)) {
+            return !__register(__joinOp);
+          }
+        }
+      } while (true);
+    }
+
+    bool __register(__counting_scopes::__state_base& __joinOp) noexcept {
+      auto* ptr = __waitingJoinOps_.load(std::memory_order_relaxed);
+
+      do {
+        if (ptr == this) {
+          // __waitingJoinOps_ == this when the list has been cleared
+          return false;
+        }
+
+        // make __joinOp's next point to the current head
+        __joinOp.__next_ = static_cast<__counting_scopes::__state_base*>(ptr);
+      } while (
+        // try to make the head point to __joinOp
+        __waitingJoinOps_.compare_exchange_weak(
+          ptr,
+          &__joinOp,
+          // I don't know what synchronization semantics we need
+          // on success, but acquire-release feels safe
+          std::memory_order_acq_rel,
+          // on failure, we'll try again so relaxed should be ok
+          std::memory_order_relaxed));
+
+      return true;
+    }
+  };
+} // namespace stdexec