select.h

Go to the documentation of this file.

// uvco (c) 2023 Lewin Bormann. See LICENSE for specific terms.
 
#pragma once
 
#include <cstddef>
#include <fmt/core.h>
#include <fmt/ranges.h>
 
#include "uvco/loop/loop.h"
#include "uvco/promise/promise.h"
 
#include <coroutine>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>
 
namespace uvco {


template <typename... Ts> class SelectSet {
public:
  using Variant = std::variant<Promise<Ts> *...>;
  using Tuple = std::tuple<Promise<Ts> *...>;
 
  explicit SelectSet(Promise<Ts> &...promises) : promises_{&promises...} {}
 
  ~SelectSet() {
    if (!resumed_) {
      std::apply(
          [](auto *...promise) { (promise->core()->resetHandle(), ...); },
          promises_);
    }
  }
 
  [[nodiscard]] bool await_ready() noexcept {
    BOOST_ASSERT_MSG(!resumed_, "A select set can only be used once");
    const bool isReady = std::apply(
        [](auto *...promise) -> bool { return (promise->ready() || ...); },
        promises_);
    resumed_ |= isReady;
    return isReady;
  }

  std::coroutine_handle<> await_suspend(std::coroutine_handle<> handle) {
    std::apply(
        [handle](auto *...promise) {
          ((!promise->core()->stale() ? promise->core()->setHandle(handle)
                                      : (void)0),
           ...);
        },
        promises_);
    suspended_ = handle;
    return Loop::getNext();
  }

  std::vector<Variant> await_resume() {
    resumed_ = true;
    std::vector<Variant> readyPromises;
    checkPromises(readyPromises);
    // Prevent the current coroutine from being resumed again, if another
    // promise in the SelectSet became ready in the meantime.
    if (suspended_ != nullptr) {
      Loop::cancel(suspended_);
      suspended_ = nullptr;
    }
    return readyPromises;
  }
 
private:
  template <size_t Ix = 0>
  void checkPromises(std::vector<Variant> &readyPromises) {
    if constexpr (Ix < sizeof...(Ts)) {
      using PromisePtr = std::tuple_element_t<Ix, Tuple>;
      PromisePtr promise = std::get<Ix>(promises_);
      if (promise->ready()) {
        readyPromises.emplace_back(std::in_place_index<Ix>, promise);
      } else {
        // Reset promises to `init` state and remove the current coroutine's
        // handle from them so they don't resume us later.
        promise->core()->resetHandle();
      }
      checkPromises<Ix + 1>(readyPromises);
    }
  }
 
  Tuple promises_;
  std::coroutine_handle<> suspended_;
  bool resumed_ = false;
};

 
} // namespace uvco