future_timeout.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 <functional>
#include <optional>
#include <type_traits>
#include <utility>
 
#include "lib/stdcompat/type_traits.h"
#include "pw_assert/assert.h"
#include "pw_async2/channel.h"
#include "pw_async2/context.h"
#include "pw_async2/future.h"
#include "pw_async2/poll.h"
#include "pw_async2/system_time_provider.h"
#include "pw_async2/time_provider.h"
#include "pw_async2/value_future.h"
#include "pw_chrono/system_clock.h"
#include "pw_function/function.h"
#include "pw_preprocessor/compiler.h"
#include "pw_status/status.h"
 
namespace pw::async2 {
 
struct EmptyReadyResolution {
  using value_type = ReadyType;
  constexpr Poll<> operator()() const { return Ready(); }
};
 
template <typename T>
struct ReadyFunctionResultResolution {
  using value_type = T;
  constexpr explicit ReadyFunctionResultResolution(
      pw::Function<T()> ready_value_fn)
      : ready_(std::move(ready_value_fn)) {}
 
  constexpr Poll<T> operator()() const {
    return Ready<T>(std::in_place_t{}, ready_());
  }
 
 private:
  pw::Function<T()> ready_;
};
 
template <typename T, typename ConstantType, typename CastType = T>
struct ReadyConstantValueResolution {
  using value_type = T;
  constexpr ReadyConstantValueResolution() = default;
  constexpr explicit ReadyConstantValueResolution(
      const cpp20::remove_cvref_t<ConstantType>& value)
      : constant_(value) {}
  constexpr explicit ReadyConstantValueResolution(
      cpp20::remove_cvref_t<ConstantType>&& value)
      : constant_(std::move(value)) {}
 
  constexpr Poll<T> operator()() const {
    return Ready<T>(std::in_place_t{}, CastType(constant_));
  }
 
 private:
  cpp20::remove_cvref_t<ConstantType> constant_;
};
 
template <typename T>
using DeadlineExceededResolution = ReadyConstantValueResolution<
    T,
    std::integral_constant<pw::Status::Code, PW_STATUS_DEADLINE_EXCEEDED>,
    pw::Status>;
 
template <
    typename T,
    typename PrimaryFuture,
    typename TimeoutFuture,
    typename TimeoutResolution,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<PrimaryFuture>>,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<TimeoutFuture>>,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<TimeoutResolution>>>
class [[nodiscard]] FutureWithTimeout
    : public Future<
          FutureWithTimeout<T, PrimaryFuture, TimeoutFuture, TimeoutResolution>,
          T> {
  static_assert(
      is_future_v<PrimaryFuture>,
      "FutureWithTimeout can only be used when PrimaryFuture is a Future type");
  static_assert(
      is_future_v<TimeoutFuture>,
      "FutureWithTimeout can only be used when TimeoutFuture is a Future type");
 
 public:
  FutureWithTimeout(PrimaryFuture&& primary_future,
                    TimeoutFuture&& timeout_future,
                    TimeoutResolution&& timeout_resolution)
      : state_{std::in_place,
               State{.primary_future = std::move(primary_future),
                     .timeout_future_ = std::move(timeout_future),
                     .timeout_resolution_ = std::move(timeout_resolution)}} {}
 
 private:
  using Base = Future<
      FutureWithTimeout<T, PrimaryFuture, TimeoutFuture, TimeoutResolution>,
      T>;
  friend Base;
 
  Poll<typename Base::value_type> DoPend(Context& cx) {
    PW_DASSERT(state_.has_value());
    auto result = state_->primary_future.Pend(cx);
    if (result.IsReady()) {
      return Ready<typename Base::value_type>(std::in_place_t{},
                                              std::move(result).value());
    }
 
    if (state_->timeout_future_.Pend(cx).IsReady()) {
      return state_->timeout_resolution_();
    }
 
    return Pending();
  }
 
  void DoMarkComplete() { state_.reset(); }
  [[nodiscard]] bool DoIsComplete() const { return !state_.has_value(); }
 
  struct State {
    PrimaryFuture primary_future;
    TimeoutFuture timeout_future_;
    PW_NO_UNIQUE_ADDRESS TimeoutResolution timeout_resolution_;
  };
  std::optional<State> state_;
};
 
template <typename T,
          typename PrimaryFuture,
          typename TimeoutFuture,
          typename TimeoutResolution>
auto CreateFutureWithTimeout(PrimaryFuture&& primary_future,
                             TimeoutFuture&& timeout_future,
                             TimeoutResolution&& timeout_resolution)
    -> FutureWithTimeout<T,
                         std::decay_t<PrimaryFuture>,
                         std::decay_t<TimeoutFuture>,
                         std::decay_t<TimeoutResolution>> {
  return FutureWithTimeout<T,
                           std::decay_t<PrimaryFuture>,
                           std::decay_t<TimeoutFuture>,
                           std::decay_t<TimeoutResolution>>(
      std::forward<PrimaryFuture>(primary_future),
      std::forward<TimeoutFuture>(timeout_future),
      std::forward<TimeoutResolution>(timeout_resolution));
}
 
template <
    typename PrimaryFuture,
    typename TimeProvider,
    typename Duration,
    int&... kExplicitGuard,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<PrimaryFuture>>>
auto Timeout(PrimaryFuture&& primary_future,
             TimeProvider& time_provider,
             Duration delay) {
  using ResultType = Result<typename std::decay_t<PrimaryFuture>::value_type>;
  return CreateFutureWithTimeout<ResultType>(
      std::forward<PrimaryFuture>(primary_future),
      time_provider.WaitFor(delay),
      DeadlineExceededResolution<ResultType>());
}
 
template <
    typename PrimaryFuture,
    int&... kExplicitGuard,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<PrimaryFuture>>>
auto Timeout(PrimaryFuture&& primary_future,
             typename chrono::SystemClock::duration delay) {
  using ResultType = Result<typename std::decay_t<PrimaryFuture>::value_type>;
  return CreateFutureWithTimeout<ResultType>(
      std::forward<PrimaryFuture>(primary_future),
      GetSystemTimeProvider().WaitFor(delay),
      DeadlineExceededResolution<ResultType>());
}
 
template <typename Clock>
void TimeoutOr([[maybe_unused]] ValueFuture<void>&& future,
               [[maybe_unused]] TimeProvider<Clock>& time_provider,
               [[maybe_unused]] typename Clock::duration delay) {
  static_assert(false, "ValueFuture<void> cannot be used with TimeoutOr");
}
 
template <typename PrimaryFuture, typename Clock, typename U>
auto TimeoutOr(PrimaryFuture&& primary_future,
               TimeProvider<Clock>& time_provider,
               typename Clock::duration delay,
               U&& value_or_fn_on_timeout) {
  using value_type = typename PrimaryFuture::value_type;
 
  static_assert(
      std::is_invocable_v<U> || std::is_convertible_v<U, value_type>,
      "value_or_fn_on_timeout (U) must be callable, or convert to the "
      "value_type of the primary_future");
 
  if constexpr (std::is_invocable_v<U>) {
    return CreateFutureWithTimeout<value_type>(
        std::forward<PrimaryFuture>(primary_future),
        time_provider.WaitFor(delay),
        ReadyFunctionResultResolution<value_type>(
            std::forward<U>(value_or_fn_on_timeout)));
  } else if constexpr (std::is_convertible_v<U, value_type>) {
    return CreateFutureWithTimeout<value_type>(
        std::forward<PrimaryFuture>(primary_future),
        time_provider.WaitFor(delay),
        ReadyConstantValueResolution<value_type, U>(
            std::forward<U>(value_or_fn_on_timeout)));
  }
}
 
template <
    typename PrimaryFuture,
    typename... Args,
    int&... kExplicitGuard,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<PrimaryFuture>>>
auto TimeoutOr(PrimaryFuture&& primary_future,
               typename chrono::SystemClock::duration delay,
               Args&&... args) {
  return TimeoutOr<typename std::decay_t<PrimaryFuture>>(
      std::forward<PrimaryFuture>(primary_future),
      GetSystemTimeProvider(),
      delay,
      std::forward<Args>(args)...);
}
 
template <typename T, typename Clock>
auto TimeoutOrClosed(ReceiveFuture<T>&& future,
                     TimeProvider<Clock>& time_provider,
                     typename Clock::duration delay) {
  return CreateFutureWithTimeout<std::optional<T>>(
      std::move(future),
      time_provider.WaitFor(delay),
      ReadyConstantValueResolution<std::optional<T>, std::nullopt_t>(
          std::nullopt));
}
 
template <typename T, typename Clock>
auto TimeoutOrClosed(SendFuture<T>&& future,
                     TimeProvider<Clock>& time_provider,
                     typename Clock::duration delay) {
  return CreateFutureWithTimeout<T>(
      std::move(future),
      time_provider.WaitFor(delay),
      ReadyConstantValueResolution<bool, std::false_type>());
}
 
template <typename T, typename Clock>
auto TimeoutOrClosed(ReserveSendFuture<T>&& future,
                     TimeProvider<Clock>& time_provider,
                     typename Clock::duration delay) {
  using ResultType = typename ReserveSendFuture<T>::value_type;
  return CreateFutureWithTimeout<ResultType>(
      std::move(future),
      time_provider.WaitFor(delay),
      ReadyConstantValueResolution<std::optional<SendReservation<T>>,
                                   std::nullopt_t>(std::nullopt));
}
 
template <
    typename PrimaryFuture,
    int&... kExplicitGuard,
    typename = std::enable_if_t<!std::is_lvalue_reference_v<PrimaryFuture>>>
auto TimeoutOrClosed(PrimaryFuture&& primary_future,
                     typename chrono::SystemClock::duration delay) {
  return TimeoutOrClosed(std::forward<PrimaryFuture>(primary_future),
                         GetSystemTimeProvider(),
                         delay);
}
 
template <typename T, typename Clock = chrono::SystemClock>
using ValueFutureWithTimeout =
    decltype(Timeout(std::declval<ValueFuture<T>&&>(),
                     std::declval<TimeProvider<Clock>&>(),
                     std::declval<typename Clock::duration>()));
 
template <typename T, typename U, typename Clock = chrono::SystemClock>
using ValueFutureWithTimeoutOr =
    decltype(TimeoutOr(std::declval<ValueFuture<T>&&>(),
                       std::declval<TimeProvider<Clock>&>(),
                       std::declval<typename Clock::duration>(),
                       std::declval<U&&>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using SendFutureWithTimeout =
    decltype(Timeout(std::declval<SendFuture<T>&&>(),
                     std::declval<TimeProvider<Clock>&>(),
                     std::declval<typename Clock::duration>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using ReceiveFutureWithTimeout =
    decltype(Timeout(std::declval<ReceiveFuture<T>&&>(),
                     std::declval<TimeProvider<Clock>&>(),
                     std::declval<typename Clock::duration>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using ReserveSendFutureWithTimeout =
    decltype(Timeout(std::declval<ReserveSendFuture<T>&&>(),
                     std::declval<TimeProvider<Clock>&>(),
                     std::declval<typename Clock::duration>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using SendFutureWithTimeoutOrClosed =
    decltype(TimeoutOrClosed(std::declval<SendFuture<T>&&>(),
                             std::declval<TimeProvider<Clock>&>(),
                             std::declval<typename Clock::duration>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using ReceiveFutureWithTimeoutOrClosed =
    decltype(TimeoutOrClosed(std::declval<ReceiveFuture<T>&&>(),
                             std::declval<TimeProvider<Clock>&>(),
                             std::declval<typename Clock::duration>()));
 
template <typename T, typename Clock = chrono::SystemClock>
using ReserveSendFutureWithTimeoutOrClosed =
    decltype(TimeoutOrClosed(std::declval<ReserveSendFuture<T>&&>(),
                             std::declval<TimeProvider<Clock>&>(),
                             std::declval<typename Clock::duration>()));
 
template <typename FutureType, typename Clock = chrono::SystemClock>
using CustomFutureWithTimeout =
    decltype(Timeout(std::declval<FutureType&&>(),
                     std::declval<TimeProvider<Clock>&>(),
                     std::declval<typename Clock::duration>()));
 
template <typename FutureType, typename U, typename Clock = chrono::SystemClock>
using CustomFutureWithTimeoutOr =
    decltype(TimeoutOr(std::declval<FutureType&&>(),
                       std::declval<TimeProvider<Clock>&>(),
                       std::declval<typename Clock::duration>(),
                       std::declval<U&&>()));
 
}  // namespace pw::async2