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"
 
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 pw::async2