irange.h 2.9 KB

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