ArrayRef.h 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383
  1. //===--- ArrayRef.h - Array Reference Wrapper -------------------*- C++ -*-===//
  2. //
  3. // The LLVM Compiler Infrastructure
  4. //
  5. // This file is distributed under the University of Illinois Open Source
  6. // License. See LICENSE.TXT for details.
  7. //
  8. //===----------------------------------------------------------------------===//
  9. // ATen: modified from llvm::ArrayRef.
  10. // removed llvm-specific functionality
  11. // removed some implicit const -> non-const conversions that rely on
  12. // complicated std::enable_if meta-programming
  13. // removed a bunch of slice variants for simplicity...
  14. #pragma once
  15. #include <c10/macros/Macros.h>
  16. #include <c10/util/Exception.h>
  17. #include <c10/util/SmallVector.h>
  18. #include <array>
  19. #include <cstddef>
  20. #include <cstdint>
  21. #include <initializer_list>
  22. #include <iterator>
  23. #include <ostream>
  24. #include <type_traits>
  25. #include <vector>
  26. namespace c10 {
  27. /// ArrayRef - Represent a constant reference to an array (0 or more elements
  28. /// consecutively in memory), i.e. a start pointer and a length. It allows
  29. /// various APIs to take consecutive elements easily and conveniently.
  30. ///
  31. /// This class does not own the underlying data, it is expected to be used in
  32. /// situations where the data resides in some other buffer, whose lifetime
  33. /// extends past that of the ArrayRef. For this reason, it is not in general
  34. /// safe to store an ArrayRef.
  35. ///
  36. /// This is intended to be trivially copyable, so it should be passed by
  37. /// value.
  38. template <typename T>
  39. class ArrayRef final {
  40. public:
  41. using iterator = const T*;
  42. using const_iterator = const T*;
  43. using size_type = size_t;
  44. using value_type = T;
  45. using reverse_iterator = std::reverse_iterator<iterator>;
  46. private:
  47. /// The start of the array, in an external buffer.
  48. const T* Data;
  49. /// The number of elements.
  50. size_type Length;
  51. void debugCheckNullptrInvariant() {
  52. TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
  53. Data != nullptr || Length == 0,
  54. "created ArrayRef with nullptr and non-zero length! std::optional relies on this being illegal");
  55. }
  56. public:
  57. /// @name Constructors
  58. /// @{
  59. /// Construct an empty ArrayRef.
  60. /* implicit */ constexpr ArrayRef() : Data(nullptr), Length(0) {}
  61. /// Construct an ArrayRef from a single element.
  62. // TODO Make this explicit
  63. constexpr ArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
  64. /// Construct an ArrayRef from a pointer and length.
  65. constexpr ArrayRef(const T* data, size_t length)
  66. : Data(data), Length(length) {
  67. debugCheckNullptrInvariant();
  68. }
  69. /// Construct an ArrayRef from a range.
  70. constexpr ArrayRef(const T* begin, const T* end)
  71. : Data(begin), Length(end - begin) {
  72. debugCheckNullptrInvariant();
  73. }
  74. /// Construct an ArrayRef from a SmallVector. This is templated in order to
  75. /// avoid instantiating SmallVectorTemplateCommon<T> whenever we
  76. /// copy-construct an ArrayRef.
  77. template <typename U>
  78. /* implicit */ ArrayRef(const SmallVectorTemplateCommon<T, U>& Vec)
  79. : Data(Vec.data()), Length(Vec.size()) {
  80. debugCheckNullptrInvariant();
  81. }
  82. template <
  83. typename Container,
  84. typename U = decltype(std::declval<Container>().data()),
  85. typename = std::enable_if_t<
  86. (std::is_same_v<U, T*> || std::is_same_v<U, T const*>)>>
  87. /* implicit */ ArrayRef(const Container& container)
  88. : Data(container.data()), Length(container.size()) {
  89. debugCheckNullptrInvariant();
  90. }
  91. /// Construct an ArrayRef from a std::vector.
  92. // The enable_if stuff here makes sure that this isn't used for
  93. // std::vector<bool>, because ArrayRef can't work on a std::vector<bool>
  94. // bitfield.
  95. template <typename A>
  96. /* implicit */ ArrayRef(const std::vector<T, A>& Vec)
  97. : Data(Vec.data()), Length(Vec.size()) {
  98. static_assert(
  99. !std::is_same_v<T, bool>,
  100. "ArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
  101. }
  102. /// Construct an ArrayRef from a std::array
  103. template <size_t N>
  104. /* implicit */ constexpr ArrayRef(const std::array<T, N>& Arr)
  105. : Data(Arr.data()), Length(N) {}
  106. /// Construct an ArrayRef from a C array.
  107. template <size_t N>
  108. // NOLINTNEXTLINE(*c-arrays*)
  109. /* implicit */ constexpr ArrayRef(const T (&Arr)[N]) : Data(Arr), Length(N) {}
  110. /// Construct an ArrayRef from a std::initializer_list.
  111. /* implicit */ constexpr ArrayRef(const std::initializer_list<T>& Vec)
  112. : Data(
  113. std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
  114. : std::begin(Vec)),
  115. Length(Vec.size()) {}
  116. /// @}
  117. /// @name Simple Operations
  118. /// @{
  119. constexpr iterator begin() const {
  120. return Data;
  121. }
  122. constexpr iterator end() const {
  123. return Data + Length;
  124. }
  125. // These are actually the same as iterator, since ArrayRef only
  126. // gives you const iterators.
  127. constexpr const_iterator cbegin() const {
  128. return Data;
  129. }
  130. constexpr const_iterator cend() const {
  131. return Data + Length;
  132. }
  133. constexpr reverse_iterator rbegin() const {
  134. return reverse_iterator(end());
  135. }
  136. constexpr reverse_iterator rend() const {
  137. return reverse_iterator(begin());
  138. }
  139. /// Check if all elements in the array satisfy the given expression
  140. constexpr bool allMatch(const std::function<bool(const T&)>& pred) const {
  141. return std::all_of(cbegin(), cend(), pred);
  142. }
  143. /// empty - Check if the array is empty.
  144. constexpr bool empty() const {
  145. return Length == 0;
  146. }
  147. constexpr const T* data() const {
  148. return Data;
  149. }
  150. /// size - Get the array size.
  151. constexpr size_t size() const {
  152. return Length;
  153. }
  154. /// front - Get the first element.
  155. constexpr const T& front() const {
  156. TORCH_CHECK(
  157. !empty(), "ArrayRef: attempted to access front() of empty list");
  158. return Data[0];
  159. }
  160. /// back - Get the last element.
  161. constexpr const T& back() const {
  162. TORCH_CHECK(!empty(), "ArrayRef: attempted to access back() of empty list");
  163. return Data[Length - 1];
  164. }
  165. /// equals - Check for element-wise equality.
  166. constexpr bool equals(ArrayRef RHS) const {
  167. return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
  168. }
  169. /// slice(n, m) - Take M elements of the array starting at element N
  170. constexpr ArrayRef<T> slice(size_t N, size_t M) const {
  171. TORCH_CHECK(
  172. N + M <= size(),
  173. "ArrayRef: invalid slice, N = ",
  174. N,
  175. "; M = ",
  176. M,
  177. "; size = ",
  178. size());
  179. return ArrayRef<T>(data() + N, M);
  180. }
  181. /// slice(n) - Chop off the first N elements of the array.
  182. constexpr ArrayRef<T> slice(size_t N) const {
  183. TORCH_CHECK(
  184. N <= size(), "ArrayRef: invalid slice, N = ", N, "; size = ", size());
  185. return slice(N, size() - N);
  186. }
  187. /// @}
  188. /// @name Operator Overloads
  189. /// @{
  190. constexpr const T& operator[](size_t Index) const {
  191. return Data[Index];
  192. }
  193. /// Vector compatibility
  194. constexpr const T& at(size_t Index) const {
  195. TORCH_CHECK(
  196. Index < Length,
  197. "ArrayRef: invalid index Index = ",
  198. Index,
  199. "; Length = ",
  200. Length);
  201. return Data[Index];
  202. }
  203. /// Disallow accidental assignment from a temporary.
  204. ///
  205. /// The declaration here is extra complicated so that "arrayRef = {}"
  206. /// continues to select the move assignment operator.
  207. template <typename U>
  208. std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
  209. // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
  210. U&& Temporary) = delete;
  211. /// Disallow accidental assignment from a temporary.
  212. ///
  213. /// The declaration here is extra complicated so that "arrayRef = {}"
  214. /// continues to select the move assignment operator.
  215. template <typename U>
  216. std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
  217. std::initializer_list<U>) = delete;
  218. /// @}
  219. /// @name Expensive Operations
  220. /// @{
  221. std::vector<T> vec() const {
  222. return std::vector<T>(Data, Data + Length);
  223. }
  224. /// @}
  225. };
  226. template <typename T>
  227. std::ostream& operator<<(std::ostream& out, ArrayRef<T> list) {
  228. int i = 0;
  229. out << "[";
  230. for (const auto& e : list) {
  231. if (i++ > 0)
  232. out << ", ";
  233. out << e;
  234. }
  235. out << "]";
  236. return out;
  237. }
  238. /// @name ArrayRef Convenience constructors
  239. /// @{
  240. /// Construct an ArrayRef from a single element.
  241. template <typename T>
  242. ArrayRef<T> makeArrayRef(const T& OneElt) {
  243. return OneElt;
  244. }
  245. /// Construct an ArrayRef from a pointer and length.
  246. template <typename T>
  247. ArrayRef<T> makeArrayRef(const T* data, size_t length) {
  248. return ArrayRef<T>(data, length);
  249. }
  250. /// Construct an ArrayRef from a range.
  251. template <typename T>
  252. ArrayRef<T> makeArrayRef(const T* begin, const T* end) {
  253. return ArrayRef<T>(begin, end);
  254. }
  255. /// Construct an ArrayRef from a SmallVector.
  256. template <typename T>
  257. ArrayRef<T> makeArrayRef(const SmallVectorImpl<T>& Vec) {
  258. return Vec;
  259. }
  260. /// Construct an ArrayRef from a SmallVector.
  261. template <typename T, unsigned N>
  262. ArrayRef<T> makeArrayRef(const SmallVector<T, N>& Vec) {
  263. return Vec;
  264. }
  265. /// Construct an ArrayRef from a std::vector.
  266. template <typename T>
  267. ArrayRef<T> makeArrayRef(const std::vector<T>& Vec) {
  268. return Vec;
  269. }
  270. /// Construct an ArrayRef from a std::array.
  271. template <typename T, std::size_t N>
  272. ArrayRef<T> makeArrayRef(const std::array<T, N>& Arr) {
  273. return Arr;
  274. }
  275. /// Construct an ArrayRef from an ArrayRef (no-op) (const)
  276. template <typename T>
  277. ArrayRef<T> makeArrayRef(const ArrayRef<T>& Vec) {
  278. return Vec;
  279. }
  280. /// Construct an ArrayRef from an ArrayRef (no-op)
  281. template <typename T>
  282. ArrayRef<T>& makeArrayRef(ArrayRef<T>& Vec) {
  283. return Vec;
  284. }
  285. /// Construct an ArrayRef from a C array.
  286. template <typename T, size_t N>
  287. // NOLINTNEXTLINE(*c-arrays*)
  288. ArrayRef<T> makeArrayRef(const T (&Arr)[N]) {
  289. return ArrayRef<T>(Arr);
  290. }
  291. // WARNING: Template instantiation will NOT be willing to do an implicit
  292. // conversions to get you to an c10::ArrayRef, which is why we need so
  293. // many overloads.
  294. template <typename T>
  295. bool operator==(c10::ArrayRef<T> a1, c10::ArrayRef<T> a2) {
  296. return a1.equals(a2);
  297. }
  298. template <typename T>
  299. bool operator!=(c10::ArrayRef<T> a1, c10::ArrayRef<T> a2) {
  300. return !a1.equals(a2);
  301. }
  302. template <typename T>
  303. bool operator==(const std::vector<T>& a1, c10::ArrayRef<T> a2) {
  304. return c10::ArrayRef<T>(a1).equals(a2);
  305. }
  306. template <typename T>
  307. bool operator!=(const std::vector<T>& a1, c10::ArrayRef<T> a2) {
  308. return !c10::ArrayRef<T>(a1).equals(a2);
  309. }
  310. template <typename T>
  311. bool operator==(c10::ArrayRef<T> a1, const std::vector<T>& a2) {
  312. return a1.equals(c10::ArrayRef<T>(a2));
  313. }
  314. template <typename T>
  315. bool operator!=(c10::ArrayRef<T> a1, const std::vector<T>& a2) {
  316. return !a1.equals(c10::ArrayRef<T>(a2));
  317. }
  318. using IntArrayRef = ArrayRef<int64_t>;
  319. using IntList [[deprecated(
  320. "This alias is deprecated because it doesn't make ownership semantics obvious. Use IntArrayRef instead!")]] =
  321. ArrayRef<int64_t>;
  322. } // namespace c10