select.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 <algorithm>
#include <tuple>
#include <utility>
#include <variant>
 
#include "pw_async2/dispatcher.h"
#include "pw_async2/future.h"
#include "pw_containers/optional_tuple.h"
 
namespace pw::async2 {
 
 
template <typename... Pendables>
class Selector;
 
namespace experimental {
 
template <typename... Futures>
class SelectFuture;
 
}  // namespace experimental
 
template <size_t I, typename T>
struct SelectResult {
  static constexpr size_t kIndex = I;
  T value;
 
 private:
  template <typename... Pendables>
  friend class Selector;
  template <typename... Futures>
  friend class experimental::SelectFuture;
 
  explicit SelectResult(T&& val) : value(std::move(val)) {}
};
 
struct AllPendablesCompleted {
  static constexpr size_t kIndex = std::numeric_limits<size_t>::max();
  std::nullopt_t value = std::nullopt;
};
 
 
namespace internal {
 
template <typename Is, typename... Pendables>
struct SelectVariantType;
 
template <size_t... Is, typename... Pendables>
struct SelectVariantType<std::index_sequence<Is...>, Pendables...> {
  using type = std::variant<SelectResult<Is, PendOutputOf<Pendables>>...,
                            AllPendablesCompleted>;
};
 
template <typename ResultVariant,
          typename AllPendablesCompletedHandler,
          typename ReadyHandlers,
          std::size_t... Is>
void VisitSelectResult(ResultVariant&& variant,
                       AllPendablesCompletedHandler&& all_completed,
                       ReadyHandlers&& ready_handlers,
                       std::index_sequence<Is...>) {
  std::visit(
      [&](auto&& arg) {
        using ArgType = std::decay_t<decltype(arg)>;
        if constexpr (std::is_same_v<ArgType, AllPendablesCompleted>) {
          std::invoke(std::forward<AllPendablesCompletedHandler>(all_completed),
                      AllPendablesCompleted{});
        } else {
          std::invoke(std::get<ArgType::kIndex>(
                          std::forward<ReadyHandlers>(ready_handlers)),
                      std::move(arg.value));
        }
      },
      std::forward<ResultVariant>(variant));
}
 
}  // namespace internal
 
 
template <typename... Pendables>
class Selector {
 public:
  static_assert(sizeof...(Pendables) > 0,
                "Cannot select over an empty set of pendables");
 
  using ResultVariant = typename internal::SelectVariantType<
      std::make_index_sequence<sizeof...(Pendables)>,
      Pendables...>::type;
 
 public:
  explicit Selector(Pendables&&... pendables)
      : pendables_(std::move(pendables)...) {}
 
  Poll<ResultVariant> Pend(Context& cx) { return PendFrom<0>(cx, false); }
 
 private:
  template <size_t kTupleIndex>
  Poll<ResultVariant> PendFrom(Context& cx, bool has_active_pendable) {
    if constexpr (kTupleIndex < sizeof...(Pendables)) {
      auto& pendable = std::get<kTupleIndex>(pendables_);
 
      if (pendable.completed()) {
        return PendFrom<kTupleIndex + 1>(cx, has_active_pendable);
      }
 
      using value_type =
          PendOutputOf<std::tuple_element_t<kTupleIndex, decltype(pendables_)>>;
      Poll<value_type> result = pendable.Pend(cx);
 
      if (result.IsReady()) {
        return Ready(ResultVariant(
            std::in_place_index<kTupleIndex>,
            SelectResult<kTupleIndex, value_type>(std::move(*result))));
      }
      return PendFrom<kTupleIndex + 1>(cx, true);
    } else {
      // All pendables have been polled and none were ready.
 
      if (!has_active_pendable) {
        return Ready(AllPendablesCompleted{});
      }
      return Pending();
    }
  }
 
  std::tuple<Pendables...> pendables_;
};
 
template <typename... Pendables>
Selector(Pendables&&...) -> Selector<Pendables...>;
 
template <typename... Pendables>
auto Select(Context& cx, Pendables&&... pendables) {
  Selector selector(std::forward<Pendables>(pendables)...);
  return selector.Pend(cx);
}
 
template <typename ResultVariant,
          typename AllPendablesCompletedHandler,
          typename... ReadyHandler>
void VisitSelectResult(
    ResultVariant&& variant,
    AllPendablesCompletedHandler&& on_all_pendables_completed,
    ReadyHandler&&... on_ready) {
  static_assert(std::variant_size_v<std::decay_t<ResultVariant>> - 1 ==
                    sizeof...(ReadyHandler),
                "Number of handlers must match the number of pendables.");
  internal::VisitSelectResult(
      std::forward<ResultVariant>(variant),
      std::forward<AllPendablesCompletedHandler>(on_all_pendables_completed),
      std::forward_as_tuple(std::forward<ReadyHandler>(on_ready)...),
      std::make_index_sequence<sizeof...(ReadyHandler)>{});
}
 
 
namespace experimental {
 
template <typename... Futures>
class SelectFuture
    : public Future<SelectFuture<Futures...>,
                    OptionalTuple<typename Futures::value_type...>> {
 public:
  static_assert(sizeof...(Futures) > 0,
                "Cannot select over an empty set of futures");
 
  using ResultTuple = OptionalTuple<typename Futures::value_type...>;
  using Base = Future<SelectFuture<Futures...>, ResultTuple>;
 
  explicit SelectFuture(Futures&&... futures)
      : futures_(std::move(futures)...) {}
 
 private:
  friend Base;
 
  static constexpr auto kTupleIndexSequence =
      std::make_index_sequence<sizeof...(Futures)>();
 
  Poll<ResultTuple> DoPend(Context& cx) {
    ResultTuple tuple;
    PendAll(cx, kTupleIndexSequence, tuple);
    if (!tuple.empty()) {
      return tuple;
    }
    return Pending();
  }
 
  void DoMarkComplete() {}
  bool DoIsComplete() const { return IsComplete<0>(); }
 
  template <size_t kTupleIndex>
  bool IsComplete() const {
    if constexpr (kTupleIndex < sizeof...(Futures)) {
      auto& future = std::get<kTupleIndex>(futures_);
      if (future.is_complete()) {
        return true;
      }
      return IsComplete<kTupleIndex + 1>();
    } else {
      return false;
    }
  }
 
  template <size_t... Is>
  void PendAll(Context& cx, std::index_sequence<Is...>, ResultTuple& result) {
    (PendFuture<Is>(cx, result), ...);
  }
 
  template <size_t kTupleIndex>
  void PendFuture(Context& cx, ResultTuple& result) {
    auto& future = std::get<kTupleIndex>(futures_);
    if (future.is_complete()) {
      return;
    }
 
    auto poll = future.Pend(cx);
    if (poll.IsReady()) {
      result.template emplace<kTupleIndex>(std::move(*poll));
    }
  }
 
  std::tuple<Futures...> futures_;
};
 
template <typename... Futures>
SelectFuture(Futures&&...) -> SelectFuture<Futures...>;
 
template <typename... Futures>
SelectFuture<Futures...> Select(Futures&&... futures) {
  static_assert(std::conjunction_v<is_future<Futures>...>,
                "All arguments to Select must be Future types");
  return SelectFuture(std::forward<Futures>(futures)...);
}
 
}  // namespace experimental
 
}  // namespace pw::async2