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"
 
namespace pw::async2 {
 
template <typename... Pendables>
class Selector;
 
template <size_t I, typename T>
struct SelectResult {
  static constexpr size_t kIndex = I;
  T value;
 
 private:
  template <typename... Pendables>
  friend class Selector;
 
  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 OutputType =
          PendOutputOf<std::tuple_element_t<kTupleIndex, decltype(pendables_)>>;
      Poll<OutputType> result = pendable.Pend(cx);
 
      if (result.IsReady()) {
        return Ready(ResultVariant(
            std::in_place_index<kTupleIndex>,
            SelectResult<kTupleIndex, OutputType>(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 pw::async2