future.h

// Copyright 2025 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// 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 <mutex>
#include <optional>
#include <type_traits>
 
#include "pw_assert/assert.h"
#include "pw_async2/context.h"
#include "pw_async2/poll.h"
#include "pw_containers/intrusive_list.h"
#include "pw_memory/container_of.h"
#include "pw_sync/interrupt_spin_lock.h"
 
namespace pw::async2 {
 
 
template <typename Derived, typename T>
class Future {
 public:
  using value_type = std::conditional_t<std::is_void_v<T>, ReadyType, T>;
 
  Poll<value_type> Pend(Context& cx) {
    PW_ASSERT(!is_complete());
    Poll<value_type> poll = derived().DoPend(cx);
    if (poll.IsReady()) {
      derived().DoMarkComplete();
    }
    return poll;
  }
 
  bool is_complete() const { return derived().DoIsComplete(); }
 
 protected:
  constexpr Future() = default;
 
 private:
  Derived& derived() { return static_cast<Derived&>(*this); }
  const Derived& derived() const { return static_cast<const Derived&>(*this); }
};
 
template <typename T, typename = void>
struct is_future : std::false_type {};
 
template <typename T>
struct is_future<
    T,
    std::void_t<typename T::value_type,
                decltype(std::declval<T&>().Pend(std::declval<Context&>())),
                decltype(std::declval<const T&>().is_complete())>>
    : std::is_convertible<decltype(&T::Pend),
                          Poll<typename T::value_type> (T::*)(Context&)> {};
 
template <typename T>
constexpr bool is_future_v =
    is_future<std::remove_cv_t<std::remove_reference_t<T>>>::value;
 
template <typename FutureType, typename Lock = sync::InterruptSpinLock>
class ListFutureProvider {
 public:
  constexpr ListFutureProvider() = default;
 
  ListFutureProvider(const ListFutureProvider&) = delete;
  ListFutureProvider& operator=(const ListFutureProvider&) = delete;
 
  void Push(FutureType& future) {
    std::lock_guard lock(lock_);
    futures_.push_back(future);
  }
 
  std::optional<std::reference_wrapper<FutureType>> Pop() {
    std::lock_guard lock(lock_);
    if (futures_.empty()) {
      return std::nullopt;
    }
    FutureType& future = futures_.front();
    futures_.pop_front();
    return std::ref(future);
  }
 
  bool empty() {
    std::lock_guard lock(lock_);
    return futures_.empty();
  }
 
  Lock& lock() { return lock_; }
 
 private:
  friend FutureType;
 
  template <typename, typename>
  friend class ListableFutureWithWaker;
 
  using LockType = Lock;
 
  IntrusiveList<FutureType> futures_;
  Lock lock_;
};
 
template <typename FutureType>
class SingleFutureProvider {
 public:
  constexpr SingleFutureProvider() = default;
 
  SingleFutureProvider(const SingleFutureProvider&) = delete;
  SingleFutureProvider& operator=(const SingleFutureProvider&) = delete;
 
  void Set(FutureType& future) {
    PW_ASSERT(!has_future());
    inner_.Push(future);
  }
 
  bool TrySet(FutureType& future) {
    if (has_future()) {
      return false;
    }
    inner_.Push(future);
    return true;
  }
 
  [[nodiscard]] std::optional<std::reference_wrapper<FutureType>> Take() {
    return inner_.Pop();
  }
 
  bool has_future() { return !inner_.empty(); }
 
 private:
  template <typename, typename>
  friend class ListableFutureWithWaker;
  friend FutureType;
 
  ListFutureProvider<FutureType> inner_;
};
 
template <typename Derived, typename T>
class ListableFutureWithWaker
    : public Future<ListableFutureWithWaker<Derived, T>, T>,
      public IntrusiveList<Derived>::Item {
 public:
  ListableFutureWithWaker(const ListableFutureWithWaker&) = delete;
  ListableFutureWithWaker& operator=(const ListableFutureWithWaker&) = delete;
 
 protected:
  class PW_LOCKABLE("pw::async2::ListableFutureWithWaker::Lock") Provider {
   public:
    void lock() PW_EXCLUSIVE_LOCK_FUNCTION() {
      if (provider_ != nullptr) {
        provider_->lock().lock();
      }
    }
 
    void unlock() PW_UNLOCK_FUNCTION() {
      if (provider_ != nullptr) {
        provider_->lock().unlock();
      }
    }
 
    Provider& operator=(ListFutureProvider<Derived>* provider) {
      provider_ = provider;
      return *this;
    }
 
    ListFutureProvider<Derived>* get() const { return provider_; }
    ListFutureProvider<Derived>& operator*() const {
      PW_ASSERT(provider_ != nullptr);
      return *provider_;
    }
    ListFutureProvider<Derived>* operator->() const {
      PW_ASSERT(provider_ != nullptr);
      return provider_;
    }
 
    explicit operator bool() const { return provider_ != nullptr; }
 
   private:
    friend class ListableFutureWithWaker<Derived, T>;
 
    explicit Provider(ListFutureProvider<Derived>* provider)
        : provider_(provider) {}
 
    ListFutureProvider<Derived>* provider_;
  };
 
  using Lock = Provider;
 
  enum ConstructedState { kMovedFrom, kReadyForCompletion };
 
  explicit ListableFutureWithWaker(ConstructedState state)
      : provider_(nullptr), complete_(state == kMovedFrom) {}
 
  explicit ListableFutureWithWaker(ListFutureProvider<Derived>& provider)
      : provider_(&provider) {
    provider.Push(derived());
  }
  explicit ListableFutureWithWaker(SingleFutureProvider<Derived>& single)
      : ListableFutureWithWaker(single.inner_) {}
 
  ~ListableFutureWithWaker() {
    if (!provider_) {
      return;
    }
    std::lock_guard guard(lock());
    if (!this->unlisted()) {
      this->unlist();
    }
  }
 
  void MoveFrom(ListableFutureWithWaker& other) {
    complete_ = std::exchange(other.complete_, true);
    provider_ = std::exchange(other.provider_, nullptr);
    waker_ = std::move(other.waker_);
 
    if (provider_) {
      std::lock_guard guard(lock());
      if (!other.unlisted()) {
        this->replace(other);
      }
    }
  }
 
  ListFutureProvider<Derived>& provider() { return *provider_; }
 
  Lock& lock() { return provider_; }
 
  void Wake() { waker_.Wake(); }
 
 private:
  using Base = Future<ListableFutureWithWaker<Derived, T>, T>;
 
  friend Base;
  friend ListFutureProvider<Derived>;
 
  Poll<typename Base::value_type> DoPend(Context& cx) {
    static_assert(
        std::is_same_v<std::remove_extent_t<decltype(Derived::kWaitReason)>,
                       const char>,
        "kWaitReason must be a character array");
 
    Poll<typename Base::value_type> poll = derived().DoPend(cx);
    if (poll.IsPending()) {
      PW_ASYNC_STORE_WAKER(cx, waker_, Derived::kWaitReason);
      Relist();
    }
    return poll;
  }
 
  void Relist() {
    std::lock_guard guard(lock());
    if (provider_ && this->unlisted()) {
      provider_->futures_.push_back(derived());
    }
  }
 
  void DoMarkComplete() { complete_ = true; }
  bool DoIsComplete() const { return complete_; }
 
  Derived& derived() { return static_cast<Derived&>(*this); }
 
  Provider provider_;
  Waker waker_;
  bool complete_ = false;
};
 
class FutureCore : public IntrusiveForwardList<FutureCore>::Item {
 public:
  constexpr FutureCore() : state_(State::kNull) {}
 
  FutureCore(FutureCore&& other) noexcept;
 
  FutureCore& operator=(FutureCore&& other) noexcept;
 
  enum ReadyForCompletion { kReadyForCompletion };
 
  enum Pending { kPending };
 
  explicit constexpr FutureCore(Pending) : state_(State::kPending) {}
 
  explicit constexpr FutureCore(ReadyForCompletion) : state_(State::kReady) {}
 
  FutureCore(const FutureCore&) = delete;
  FutureCore& operator=(const FutureCore&) = delete;
 
  ~FutureCore() = default;
 
  [[nodiscard]] bool is_pendable() const { return state_ > State::kComplete; }
 
  [[nodiscard]] bool is_complete() const { return state_ == State::kComplete; }
 
  [[nodiscard]] bool is_ready() const { return state_ == State::kReady; }
 
  void Wake() { waker_.Wake(); }
 
  void WakeAndMarkReady() {
    Wake();
    state_ = State::kReady;
  }
 
  void Unlist() { unlist(); }
 
  template <typename FutureType>
  auto DoPend(FutureType& future, Context& cx) PW_NO_LOCK_SAFETY_ANALYSIS {
    PW_ASSERT(is_pendable());
 
    auto poll = future.DoPend(cx);
    if (poll.IsPending()) {
      PW_ASYNC_STORE_WAKER(cx, waker_, FutureType::kWaitReason);
    } else {
      state_ = State::kComplete;
    }
 
    return poll;
  }
 
  [[nodiscard]] bool in_list() const { return !unlisted(); }
 
 private:
  friend class BaseFutureList;  // for IntrusiveForwardList::Item functions
 
  [[nodiscard]] bool is_initialized() const { return state_ != State::kNull; }
 
  Waker waker_;
 
  enum class State : unsigned char {
    kNull,
 
    kComplete,
 
    kReady,
 
    kPending,
  } state_;
};
 
class BaseFutureList {
 public:
  constexpr BaseFutureList() = default;
 
  bool empty() const { return list_.empty(); }
 
  void Push(FutureCore& future) { containers::PushBackSlow(list_, future); }
 
  void PushRequireEmpty(FutureCore& future);
 
  bool PushIfEmpty(FutureCore& future);
 
  FutureCore* PopIfAvailable() { return list_.empty() ? nullptr : &Pop(); }
 
  FutureCore& Pop() {
    FutureCore& future = list_.front();
    list_.pop_front();
    return future;
  }
 
  void ResolveOne() {
    list_.front().WakeAndMarkReady();
    list_.pop_front();
  }
 
  void ResolveOneIfAvailable() {
    if (!list_.empty()) {
      ResolveOne();
    }
  }
 
  void ResolveAll();
 
 protected:
  IntrusiveForwardList<FutureCore>& list() { return list_; }
 
 private:
  IntrusiveForwardList<FutureCore> list_;
};
 
template <auto kGetFutureImpl, auto kGetFutureCore>
class CustomFutureList : public BaseFutureList {
 public:
  using value_type = std::remove_reference_t<decltype(*kGetFutureImpl(
      std::declval<FutureCore*>()))>;
  using pointer = value_type*;
  using reference = value_type&;
 
  constexpr CustomFutureList() = default;
 
  void Push(FutureCore& future) { BaseFutureList::Push(future); }
  void Push(reference future) { Push(kGetFutureCore(future)); }
 
  void PushRequireEmpty(FutureCore& future) {
    BaseFutureList::PushRequireEmpty(future);
  }
  void PushRequireEmpty(reference future) {
    PushRequireEmpty(kGetFutureCore(future));
  }
 
  bool PushIfEmpty(FutureCore& future) {
    return BaseFutureList::PushIfEmpty(future);
  }
  bool PushIfEmpty(reference future) {
    return PushIfEmpty(kGetFutureCore(future));
  }
 
  pointer PopIfAvailable() {
    return kGetFutureImpl(BaseFutureList::PopIfAvailable());
  }
 
  reference Pop() { return *kGetFutureImpl(BaseFutureList::PopIfAvailable()); }
 
  template <typename Resolver>
  void ResolveAllWith(Resolver&& resolver) {
    while (!list().empty()) {
      resolver(Pop());
    }
  }
 
  template <typename Resolver>
  void ResolveOneWith(Resolver&& resolver) {
    if (!list().empty()) {
      resolver(Pop());
    }
  }
};
 
template <auto kMemberPtr>
using FutureList =
    CustomFutureList<ContainerOf<kMemberPtr, FutureCore>, MemberOf<kMemberPtr>>;
 
 
}  // namespace pw::async2