join.h

// Copyright 2024 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 "pw_async2/dispatcher.h"
#include "pw_async2/future.h"
 
namespace pw::async2 {
 
 
template <typename... Pendables>
class Join {
 private:
  static constexpr auto kTupleIndexSequence =
      std::make_index_sequence<sizeof...(Pendables)>();
  using TupleOfOutputRvalues = std::tuple<PendOutputOf<Pendables>&&...>;
 
 public:
  explicit Join(Pendables&&... pendables)
      : pendables_(std::move(pendables)...),
        outputs_(Poll<PendOutputOf<Pendables>>(Pending())...) {}
 
  Poll<TupleOfOutputRvalues> Pend(Context& cx) {
    if (!PendElements(cx, kTupleIndexSequence)) {
      return Pending();
    }
    return TakeOutputs(kTupleIndexSequence);
  }
 
 private:
  template <size_t... Is>
  bool PendElements(Context& cx, std::index_sequence<Is...>) {
    return (... && PendElement<Is>(cx));
  }
 
  template <size_t... Is>
  Poll<TupleOfOutputRvalues> TakeOutputs(std::index_sequence<Is...>) {
    return Poll<TupleOfOutputRvalues>(
        std::forward_as_tuple<PendOutputOf<Pendables>...>(TakeOutput<Is>()...));
  }
 
  template <size_t kTupleIndex>
  bool PendElement(Context& cx) {
    auto& output = std::get<kTupleIndex>(outputs_);
    if (output.IsReady()) {
      return true;
    }
    output = std::get<kTupleIndex>(pendables_).Pend(cx);
    return output.IsReady();
  }
 
  template <size_t kTupleIndex>
  PendOutputOf<typename std::tuple_element<kTupleIndex,
                                           std::tuple<Pendables...>>::type>&&
  TakeOutput() {
    return std::move(std::get<kTupleIndex>(outputs_).value());
  }
 
  std::tuple<Pendables...> pendables_;
  std::tuple<Poll<PendOutputOf<Pendables>>...> outputs_;
};
 
template <typename... Pendables>
Join(Pendables&&...) -> Join<Pendables...>;
 
 
namespace experimental {
 
template <typename... Futures>
class JoinFuture
    : public Future<JoinFuture<Futures...>,
                    std::tuple<typename Futures::value_type&&...>> {
 private:
  static constexpr auto kTupleIndexSequence =
      std::make_index_sequence<sizeof...(Futures)>();
  using TupleOfOutputRvalues = std::tuple<typename Futures::value_type&&...>;
 
  using Base = Future<JoinFuture<Futures...>, TupleOfOutputRvalues>;
  friend Base;
 
  template <typename... Fs>
  friend constexpr auto Join(Fs&&...);
 
  explicit constexpr JoinFuture(Futures&&... futures)
      : futures_(std::in_place, std::move(futures)...),
        outputs_(Poll<typename Futures::value_type>(Pending())...) {}
 
  Poll<TupleOfOutputRvalues> DoPend(Context& cx) {
    if (!PendElements(cx, kTupleIndexSequence)) {
      return Pending();
    }
    return TakeOutputs(kTupleIndexSequence);
  }
 
  void DoMarkComplete() { futures_.reset(); }
  bool DoIsComplete() const { return !futures_.has_value(); }
 
  template <size_t... Is>
  bool PendElements(Context& cx, std::index_sequence<Is...>) {
    (PendElement<Is>(cx), ...);
    return (std::get<Is>(outputs_).IsReady() && ...);
  }
 
  template <size_t kTupleIndex>
  void PendElement(Context& cx) {
    auto& output = std::get<kTupleIndex>(outputs_);
    if (!output.IsReady()) {
      output = std::get<kTupleIndex>(*futures_).Pend(cx);
    }
  }
 
  template <size_t... Is>
  Poll<TupleOfOutputRvalues> TakeOutputs(std::index_sequence<Is...>) {
    return Poll<TupleOfOutputRvalues>(
        std::forward_as_tuple<PendOutputOf<Futures>...>(TakeOutput<Is>()...));
  }
 
  template <size_t kTupleIndex>
  PendOutputOf<
      typename std::tuple_element<kTupleIndex, std::tuple<Futures...>>::type>&&
  TakeOutput() {
    return std::move(std::get<kTupleIndex>(outputs_).value());
  }
 
  static_assert(std::conjunction_v<is_future<Futures>...>,
                "All types in JoinFuture must be Future types");
  std::optional<std::tuple<Futures...>> futures_;
  std::tuple<Poll<PendOutputOf<Futures>>...> outputs_;
};
 
template <typename... Futures>
JoinFuture(Futures&&...) -> JoinFuture<Futures...>;
 
template <typename... Futures>
constexpr auto Join(Futures&&... futures) {
  static_assert(std::conjunction_v<is_future<Futures>...>,
                "All arguments to Join must be Future types");
  return JoinFuture(std::forward<Futures>(futures)...);
}
 
}  // namespace experimental
 
}  // namespace pw::async2