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 | |
12 | namespace c10 { |
13 | |
14 | namespace detail { |
15 | |
16 | template < |
17 | typename I, |
18 | bool one_sided = false, |
19 | typename std::enable_if<std::is_integral<I>::value, int>::type = 0> |
20 | struct 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 | |
70 | template < |
71 | typename I, |
72 | bool one_sided = false, |
73 | typename std::enable_if<std::is_integral<I>::value, bool>::type = true> |
74 | struct 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. |
94 | template < |
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> |
101 | integer_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 |
111 | template < |
112 | typename Integer, |
113 | typename std::enable_if<std::is_integral<Integer>::value, bool>::type = |
114 | true> |
115 | integer_range<Integer, true> irange(Integer end) { |
116 | return {Integer(), end}; |
117 | } |
118 | |
119 | } // namespace c10 |
120 | |