1// Copyright 2004-present Facebook. All Rights Reserved.
2
3#pragma once
4
5#include <c10/util/Exception.h>
6
7#include <algorithm>
8#include <iterator>
9#include <limits>
10#include <type_traits>
11
12namespace c10 {
13
14namespace detail {
15
16template <
17 typename I,
18 bool one_sided = false,
19 typename std::enable_if<std::is_integral<I>::value, int>::type = 0>
20struct integer_iterator {
21 using iterator_category = std::input_iterator_tag;
22 using value_type = I;
23 using difference_type = std::ptrdiff_t;
24 using pointer = I*;
25 using reference = I&;
26
27 explicit integer_iterator(I value) : value(value) {}
28
29 I operator*() const {
30 return value;
31 }
32
33 I const* operator->() const {
34 return &value;
35 }
36
37 integer_iterator& operator++() {
38 ++value;
39 return *this;
40 }
41
42 integer_iterator operator++(int) {
43 const auto copy = *this;
44 ++*this;
45 return copy;
46 }
47
48 bool operator==(const integer_iterator& other) const {
49 if /* constexpr -- we don't have C++17 yet, see #85969 */ (one_sided) {
50 // Range-for loops' end test is `begin != end`, not `begin <
51 // end`. To handle `c10::irange(n)` where n < 0 (which should be
52 // empty), we just make `begin != end` fail whenever `end` is
53 // negative.
54 return other.value < 0 || value == other.value;
55 } else {
56 return value == other.value;
57 }
58 }
59
60 bool operator!=(const integer_iterator& other) const {
61 return !(*this == other);
62 }
63
64 protected:
65 I value;
66};
67
68} // namespace detail
69
70template <
71 typename I,
72 bool one_sided = false,
73 typename std::enable_if<std::is_integral<I>::value, bool>::type = true>
74struct integer_range {
75 public:
76 integer_range(I begin, I end) : begin_(begin), end_(end) {}
77 using iterator = detail::integer_iterator<I, one_sided>;
78 iterator begin() const {
79 return begin_;
80 }
81 iterator end() const {
82 return end_;
83 }
84
85 private:
86 iterator begin_;
87 iterator end_;
88};
89
90/// Creates an integer range for the half-open interval [begin, end)
91/// If end<=begin, then the range is empty.
92/// The range has the type of the `end` integer; `begin` integer is
93/// cast to this type.
94template <
95 typename Integer1,
96 typename Integer2,
97 typename std::enable_if<std::is_integral<Integer1>::value, bool>::type =
98 true,
99 typename std::enable_if<std::is_integral<Integer2>::value, bool>::type =
100 true>
101integer_range<Integer2> irange(Integer1 begin, Integer2 end) {
102 // If end<=begin then the range is empty; we can achieve this effect by
103 // choosing the larger of {begin, end} as the loop terminator
104 return {
105 static_cast<Integer2>(begin),
106 std::max(static_cast<Integer2>(begin), end)};
107}
108
109/// Creates an integer range for the half-open interval [0, end)
110/// If end<=begin, then the range is empty
111template <
112 typename Integer,
113 typename std::enable_if<std::is_integral<Integer>::value, bool>::type =
114 true>
115integer_range<Integer, true> irange(Integer end) {
116 return {Integer(), end};
117}
118
119} // namespace c10
120