expanding_array.h

#pragma once

#include <c10/util/ArrayRef.h>
#include <c10/util/Exception.h>

#include <algorithm>
#include <array>
#include <cstdint>
#include <initializer_list>
#include <vector>

namespace torch {

template <size_t D, typename T = int64_t>
class ExpandingArray {
 public:
  /*implicit*/ ExpandingArray(std::initializer_list<T> list)
      : ExpandingArray(at::ArrayRef<T>(list)) {}

  /*implicit*/ ExpandingArray(at::ArrayRef<T> values) {
    // clang-format off
    AT_CHECK(
        values.size() == D,
        "Expected ", D, " values, but instead got ", values.size());
    // clang-format on
    std::copy(values.begin(), values.end(), values_.begin());
  }

  /*implicit*/ ExpandingArray(T single_size) {
    values_.fill(single_size);
  }

  /*implicit*/ ExpandingArray(const std::array<T, D>& values)
      : values_(values) {}

  std::array<T, D>& operator*() {
    return values_;
  }

  const std::array<T, D>& operator*() const {
    return values_;
  }

  std::array<T, D>* operator->() {
    return &values_;
  }

  const std::array<T, D>* operator->() const {
    return &values_;
  }

  operator at::ArrayRef<T>() const {
    return values_;
  }

  size_t size() const noexcept {
    return D;
  }

 private:
  std::array<T, D> values_;
};

template <size_t D, typename T>
std::ostream& operator<<(
    std::ostream& stream,
    const ExpandingArray<D, T>& expanding_array) {
  if (expanding_array.size() == 1) {
    return stream << expanding_array->at(0);
  }
  return stream << static_cast<at::ArrayRef<T>>(expanding_array);
}
} // namespace torch