base_rdd.h

#ifndef CPARK_BASE_RDD_H
#define CPARK_BASE_RDD_H
 
#include <concepts>
#include <ranges>
#include <utility>
#include <variant>
 
#include "cpark.h"
 
namespace cpark {
 
namespace concepts {
 
template <typename T>
concept HasId = requires(const T& t) {
  t.id();
  requires std::is_arithmetic_v<decltype(t.id())>;
};
 
template <typename T>
concept HasDependency = concepts::HasId<T> && requires(T & t) {
  { t.dependencies() } -> std::ranges::input_range;
  t.addDependency(t.id());
};
 
template <typename S>
concept Split = std::ranges::input_range<S> && concepts::HasDependency<S>;
 
template <typename R>
concept Rdd = std::ranges::input_range<R> && concepts::Split<std::ranges::range_value_t<R>>;
 
template <typename R>
concept HasBeginImpl = requires(const R& r) {
  { r.beginImpl() } -> std::input_iterator;
};
 
template <typename R>
concept HasEndImpl = requires(const R& r) {
  { r.endImpl() } -> std::sentinel_for<decltype(r.beginImpl())>;
};
 
template <class P>
concept Pair = requires(P p) {
  typename P::first_type;
  typename P::second_type;
  p.first;
  p.second;
  requires std::same_as<decltype(p.first), typename P::first_type>;
  requires std::same_as<decltype(p.second), typename P::second_type>;
};
 
template <typename R>
concept KeyValueRdd = concepts::Rdd<R> && concepts::Pair<utils::RddElementType<R>>;
 
}  // namespace concepts
 
template <typename DerivedSplit>
class BaseSplit : public std::ranges::view_interface<BaseSplit<DerivedSplit>> {
public:
  template <typename T>
  friend class BaseSplit;
 
public:
  explicit BaseSplit<DerivedSplit>(ExecutionContext* context)
      : context_{context}, split_id_{context_->getAndIncSplitId()} {}
 
  template <concepts::Split S>
  explicit BaseSplit(const BaseSplit<S>& prev)
      : context_{prev.context_}, split_id_{prev.split_id_}, dependencies_{prev.dependencies_} {}
 
  template <concepts::Split S>
  BaseSplit& operator=(const BaseSplit<S>& prev) {
    context_ = prev.context_;
    split_id_ = prev.split_id_;
    dependencies_ = prev.dependencies_;
  }
 
  template <typename T>
  BaseSplit(const BaseSplit<T>& other, bool copy_id, bool copy_dependencies)
      : context_{other.context_} {
    if (copy_id) {
      split_id_ = other.split_id_;
    } else {
      split_id_ = context_->getAndIncSplitId();
    }
    if (copy_dependencies) {
      dependencies_ = other.dependencies_;
    }
  }
 
  auto begin() const requires concepts::HasBeginImpl<DerivedSplit> {
    return static_cast<const DerivedSplit&>(*this).beginImpl();
  }
 
  auto end() const requires concepts::HasEndImpl<DerivedSplit> {
    return static_cast<const DerivedSplit&>(*this).endImpl();
  }
 
  auto dependencies() const noexcept {
    return std::ranges::subrange(std::ranges::begin(dependencies_),
                                 std::ranges::end(dependencies_));
  }
 
  void addDependency(ExecutionContext::SplitId split_id) {
    dependencies_.push_back(split_id);
    context_->markDependency(split_id_, split_id);
  }
 
  template <typename T>
  void addDependency(const BaseSplit<T>& split) {
    addDependency(split.split_id_);
  }
 
  ExecutionContext::SplitId id() const noexcept { return split_id_; }
 
protected:
  ExecutionContext* context_{};
  ExecutionContext::SplitId split_id_{};
  std::vector<ExecutionContext::SplitId> dependencies_{};
};
 
template <typename DerivedSplit, typename DerivedSplitIterator>
class CachedSplit : public BaseSplit<CachedSplit<DerivedSplit, DerivedSplitIterator>> {
public:
  template <typename T, typename U>
  friend class CachedSplit;
 
public:
  using Base = BaseSplit<CachedSplit<DerivedSplit, DerivedSplitIterator>>;
  friend Base;
  using ValueType = std::iter_value_t<DerivedSplitIterator>;
  using CacheType = std::vector<ValueType>;
 
  class Iterator : public std::random_access_iterator_tag {
  public:
    using difference_type = std::ptrdiff_t;
    using value_type = ValueType;
    using CacheIterator = std::ranges::iterator_t<const CacheType>;
    using OriginalIterator = DerivedSplitIterator;
 
    Iterator() = default;
 
    explicit Iterator(const CacheIterator& iterator) : iterator_{iterator} {
      static_assert(std::random_access_iterator<Iterator>,
                    "CachedSplit::Iterator does not satisfy random_access_iterator, please check "
                    "the implementation of the DerivedSplit.");
    }
 
    explicit Iterator(const OriginalIterator& iterator) : iterator_{iterator} {
      static_assert(std::random_access_iterator<Iterator>,
                    "CachedSplit::Iterator does not satisfy random_access_iterator, please check "
                    "the implementation of the DerivedSplit.");
    }
 
    // Member functions to implement a forward_iterator.
 
    value_type operator*() const {
      // Read the value pointed by the actual iterator inside this class.
      // First check which type of iterator is actually held by `iterator_`,
      // then get the right type from it and read the value.
      if (std::holds_alternative<CacheIterator>(iterator_)) [[unlikely]] {
        return *std::get<CacheIterator>(iterator_);
      } else {
        return *std::get<OriginalIterator>(iterator_);
      }
    }
 
    Iterator& operator++() {
      if (std::holds_alternative<CacheIterator>(iterator_)) [[unlikely]] {
        ++std::get<CacheIterator>(iterator_);
      } else {
        ++std::get<OriginalIterator>(iterator_);
      }
      return *this;
    }
 
    Iterator operator++(int) {
      auto old = *this;
      ++(*this);
      return old;
    }
 
    bool operator==(const Iterator& other) const {
      if (std::holds_alternative<CacheIterator>(iterator_) &&
          std::holds_alternative<CacheIterator>(other.iterator_)) [[unlikely]] {
        return std::get<CacheIterator>(iterator_) == std::get<CacheIterator>(other.iterator_);
      } else if (std::holds_alternative<OriginalIterator>(iterator_) &&
                 std::holds_alternative<OriginalIterator>(other.iterator_)) {
        return std::get<OriginalIterator>(iterator_) == std::get<OriginalIterator>(other.iterator_);
      } else [[unlikely]] {
        return false;
      }
    }
 
    bool operator!=(const Iterator& other) const { return !(*this == other); }
 
    // Member functions to implement bidirectional_iterator.
 
    Iterator& operator--() {
      if (std::holds_alternative<CacheIterator>(iterator_)) [[unlikely]] {
        --std::get<CacheIterator>(iterator_);
      } else {
        if constexpr (std::bidirectional_iterator<OriginalIterator>) {
          --std::get<OriginalIterator>(iterator_);
        } else {
          throw std::runtime_error("calculate-cache-on-miss not implemented yet");
        }
      }
      return *this;
    }
 
    Iterator operator--(int) {
      auto old = *this;
      ++(*this);
      return old;
    }
 
    // Member functions to implement random_access_iterator.
 
    Iterator& operator+=(const difference_type& n) {
      if (std::holds_alternative<CacheIterator>(iterator_)) [[unlikely]] {
        std::get<CacheIterator>(iterator_) += n;
      } else {
        if constexpr (std::random_access_iterator<OriginalIterator>) {
          std::get<OriginalIterator>(iterator_) += n;
        } else {
          throw std::runtime_error("Not implemented");
        }
      }
      return *this;
    }
 
    Iterator& operator-=(const difference_type& n) {
      if (std::holds_alternative<CacheIterator>(iterator_)) [[unlikely]] {
        std::get<CacheIterator>(iterator_) -= n;
      } else {
        if constexpr (std::random_access_iterator<OriginalIterator>) {
          std::get<OriginalIterator>(iterator_) -= n;
        } else {
          throw std::runtime_error("Not implemented");
        }
      }
      return *this;
    }
 
    Iterator operator+(const difference_type& n) const {
      auto res = *this;
      return res += n;
    }
 
    friend Iterator operator+(const difference_type& n, const Iterator& iter) { return iter + n; }
 
    Iterator operator-(const difference_type& n) const {
      auto res = *this;
      return res -= n;
    }
 
    value_type operator[](const difference_type& n) const { return *(*this + n); }
 
    difference_type operator-(const Iterator& other) const {
      if (std::holds_alternative<CacheIterator>(iterator_) &&
          std::holds_alternative<CacheIterator>(other.iterator_)) {
        return std::get<CacheIterator>(iterator_) - std::get<CacheIterator>(other.iterator_);
      } else if (std::holds_alternative<OriginalIterator>(iterator_) &&
                 std::holds_alternative<OriginalIterator>(other.iterator_)) {
        if constexpr (requires(const OriginalIterator& a, const OriginalIterator& b) { a - b; }) {
          return std::get<OriginalIterator>(iterator_) -
                 std::get<OriginalIterator>(other.iterator_);
        } else {
          throw std::runtime_error("Not Implemented");
        }
      } else {
        throw std::runtime_error("Bad Compare");
      }
    }
 
    // Functions to implement totally_ordered.
 
    bool operator<(const Iterator& other) const { return *this - other < 0; }
 
    bool operator>(const Iterator& other) const { return *this - other > 0; }
 
    bool operator<=(const Iterator& other) const { return *this - other <= 0; }
 
    bool operator>=(const Iterator& other) const { return *this - other >= 0; }
 
  private:
  private:
    // A variant holds either an iterator from the original split, or an iterator of the cache.
    std::variant<CacheIterator, OriginalIterator> iterator_;
  };
 
public:
  explicit CachedSplit(ExecutionContext* context) : Base{context} {
    static_assert(std::ranges::random_access_range<CachedSplit>,
                  "CachedSplit instance does not satisfy random_access_range, please check the "
                  "DerivedSplit.");
  }
 
  template <typename T, typename U>
  CachedSplit(const CachedSplit<T, U>& other, bool copy_id, bool copy_dependencies)
      : Base{other, copy_id, copy_dependencies} {
    static_assert(std::ranges::random_access_range<CachedSplit>,
                  "CachedSplit instance does not satisfy random_access_range, please check the "
                  "DerivedSplit.");
  }
 
private:
  bool shouldCache() const noexcept { return Base::context_->splitShouldCache(Base::split_id_); }
 
  bool hasCached() const noexcept { return Base::context_->splitCached(Base::split_id_); }
 
  const CacheType& getCache() const {
    return std::any_cast<const CacheType&>(Base::context_->getSplitCache(Base::split_id_));
  }
 
  void waitOrCalculate() const {
    auto future = Base::context_->template startCalculationOrGetFuture<CacheType>(
        Base::split_id_, static_cast<const DerivedSplit&>(*this).beginImpl(),
        static_cast<const DerivedSplit&>(*this).endImpl());
    future.wait();
  }
 
  auto beginImpl() const requires concepts::HasBeginImpl<DerivedSplit> {
    if (shouldCache()) [[unlikely]] {
      waitOrCalculate();
      return Iterator{std::ranges::begin(getCache())};
    } else {
      return Iterator{static_cast<const DerivedSplit&>(*this).beginImpl()};
    }
  }
 
  auto endImpl() const requires concepts::HasEndImpl<DerivedSplit> {
    if (shouldCache()) [[unlikely]] {
      waitOrCalculate();
      return Iterator{std::ranges::end(getCache())};
    } else {
      return Iterator{static_cast<const DerivedSplit&>(*this).endImpl()};
    }
  }
};
 
template <std::ranges::view V>
class ViewSplit : public CachedSplit<ViewSplit<V>, std::ranges::iterator_t<const V>> {
public:
  template <std::ranges::view T>
  friend class ViewSplit;
 
public:
  using Base = CachedSplit<ViewSplit<V>, std::ranges::iterator_t<const V>>;
  friend Base;
 
public:
  ViewSplit(V view, ExecutionContext* context) : Base{context}, view_{view} {}
 
  template <concepts::Split S>
  ViewSplit(V view, const S& prev) : Base{prev, false, false}, view_{view} {}
 
private:
  auto beginImpl() const { return std::ranges::begin(view_); }
 
  auto endImpl() const { return std::ranges::end(view_); }
 
private:
  V view_;
};
 
template <typename DerivedRdd>
class BaseRdd : public std::ranges::view_interface<BaseRdd<DerivedRdd>> {
public:
  template <typename T>
  friend class BaseRdd;
 
public:
  template <concepts::Rdd R>
  explicit BaseRdd(const BaseRdd<R>& prev, bool copy_id)
      : context_{prev.context_}, splits_num_{prev.splits_num_} {
    if (copy_id) {
      rdd_id_ = prev.rdd_id_;
    } else {
      rdd_id_ = context_->getAndIncRddId();
    }
  }
 
  template <concepts::Rdd R>
  explicit BaseRdd(const BaseRdd<R>& prev)
      : context_{prev.context_}, rdd_id_{prev.rdd_id_}, splits_num_{prev.splits_num_} {}
 
  template <concepts::Rdd R>
  BaseRdd& operator=(const BaseRdd<R>& prev) {
    context_ = prev.context_;
    rdd_id_ = prev.rdd_id_;
    splits_num_ = prev.splits_num_;
  }
 
  explicit BaseRdd(ExecutionContext* context)
      : context_{context},
        rdd_id_{context_->getAndIncRddId()},
        splits_num_{context_->getConfig().getParallelTaskNum()} {}
 
  auto begin() const requires concepts::HasBeginImpl<DerivedRdd> {
    return static_cast<const DerivedRdd&>(*this).beginImpl();
  }
 
  auto end() const requires concepts::HasEndImpl<DerivedRdd> {
    return static_cast<const DerivedRdd&>(*this).endImpl();
  }
 
  ExecutionContext::RddId id() const noexcept { return rdd_id_; }
 
protected:
  ExecutionContext* context_{};
  ExecutionContext::RddId rdd_id_{};
  size_t splits_num_{};
};
 
}  // namespace cpark
 
#endif  // CPARK_BASE_RDD_H