1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
// Copyright 2004-present Facebook. All Rights Reserved.
#pragma once
#include <c10/util/TypeSafeSignMath.h>
#include <algorithm>
#include <cstddef>
#include <iterator>
#include <type_traits>
namespace c10 {
namespace detail {
template <
typename I,
bool one_sided = false,
std::enable_if_t<std::is_integral_v<I>, int> = 0>
struct integer_iterator {
using iterator_category = std::input_iterator_tag;
using value_type = I;
using difference_type = std::ptrdiff_t;
using pointer = I*;
using reference = I&;
explicit integer_iterator(I value) : value(value) {}
I operator*() const {
return value;
}
I const* operator->() const {
return &value;
}
integer_iterator& operator++() {
++value;
return *this;
}
integer_iterator operator++(int) {
const auto copy = *this;
++*this;
return copy;
}
bool operator==(const integer_iterator& other) const {
if constexpr (one_sided) {
// Range-for loops' end test is `begin != end`, not `begin <
// end`. To handle `c10::irange(n)` where n < 0 (which should be
// empty), we just make `begin != end` fail whenever `end` is
// negative.
return is_negative(other.value) || value == other.value;
} else {
return value == other.value;
}
// Suppress "warning: missing return statement at end of non-void function"
// which Nvidia's Robert Crovella confirms is an NVCC compiler error
// here https://stackoverflow.com/a/64561686/752843 on 2020-10-27
// `__builtin_unreachable();` would be best here, but it's not
// available with all compilers. So we instead return an arbitrary
// value trusting that this line will, in fact, never be reached.
return false; // Horrible hack
}
bool operator!=(const integer_iterator& other) const {
return !(*this == other);
}
protected:
I value;
};
} // namespace detail
template <
typename I,
bool one_sided = false,
std::enable_if_t<std::is_integral_v<I>, bool> = true>
struct integer_range {
public:
integer_range(I begin, I end) : begin_(begin), end_(end) {}
using iterator = detail::integer_iterator<I, one_sided>;
iterator begin() const {
return begin_;
}
iterator end() const {
return end_;
}
private:
iterator begin_;
iterator end_;
};
/// Creates an integer range for the half-open interval [begin, end)
/// If end<=begin, then the range is empty.
/// The range has the type of the `end` integer; `begin` integer is
/// cast to this type.
template <
typename Integer1,
typename Integer2,
std::enable_if_t<std::is_integral_v<Integer1>, bool> = true,
std::enable_if_t<std::is_integral_v<Integer2>, bool> = true>
integer_range<Integer2> irange(Integer1 begin, Integer2 end) {
// If end<=begin then the range is empty; we can achieve this effect by
// choosing the larger of {begin, end} as the loop terminator
return {
static_cast<Integer2>(begin),
std::max(static_cast<Integer2>(begin), end)};
}
/// Creates an integer range for the half-open interval [0, end)
/// If end<=begin, then the range is empty
template <
typename Integer,
std::enable_if_t<std::is_integral_v<Integer>, bool> = true>
integer_range<Integer, true> irange(Integer end) {
return {Integer(), end};
}
} // namespace c10
|
