Bitset.h 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. #pragma once
  2. #include <cstddef>
  3. #if defined(_MSC_VER)
  4. #include <intrin.h>
  5. #endif
  6. namespace c10::utils {
  7. /**
  8. * This is a simple bitset class with sizeof(long long int) bits.
  9. * You can set bits, unset bits, query bits by index,
  10. * and query for the first set bit.
  11. * Before using this class, please also take a look at std::bitset,
  12. * which has more functionality and is more generic. It is probably
  13. * a better fit for your use case. The sole reason for c10::utils::bitset
  14. * to exist is that std::bitset misses a find_first_set() method.
  15. */
  16. struct bitset final {
  17. private:
  18. #if defined(_MSC_VER)
  19. // MSVCs _BitScanForward64 expects int64_t
  20. using bitset_type = int64_t;
  21. #else
  22. // POSIX ffsll expects long long int
  23. using bitset_type = long long int;
  24. #endif
  25. public:
  26. static constexpr size_t NUM_BITS() {
  27. return 8 * sizeof(bitset_type);
  28. }
  29. constexpr bitset() noexcept = default;
  30. constexpr bitset(const bitset&) noexcept = default;
  31. constexpr bitset(bitset&&) noexcept = default;
  32. // there is an issure for gcc 5.3.0 when define default function as constexpr
  33. // see https://gcc.gnu.org/bugzilla/show_bug.cgi?id=68754.
  34. bitset& operator=(const bitset&) noexcept = default;
  35. bitset& operator=(bitset&&) noexcept = default;
  36. ~bitset() = default;
  37. constexpr void set(size_t index) noexcept {
  38. bitset_ |= (static_cast<long long int>(1) << index);
  39. }
  40. constexpr void unset(size_t index) noexcept {
  41. bitset_ &= ~(static_cast<long long int>(1) << index);
  42. }
  43. constexpr bool get(size_t index) const noexcept {
  44. return bitset_ & (static_cast<long long int>(1) << index);
  45. }
  46. constexpr bool is_entirely_unset() const noexcept {
  47. return 0 == bitset_;
  48. }
  49. // Call the given functor with the index of each bit that is set
  50. template <class Func>
  51. // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
  52. void for_each_set_bit(Func&& func) const {
  53. bitset cur = *this;
  54. size_t index = cur.find_first_set();
  55. while (0 != index) {
  56. // -1 because find_first_set() is not one-indexed.
  57. index -= 1;
  58. func(index);
  59. cur.unset(index);
  60. index = cur.find_first_set();
  61. }
  62. }
  63. private:
  64. // Return the index of the first set bit. The returned index is one-indexed
  65. // (i.e. if the very first bit is set, this function returns '1'), and a
  66. // return of '0' means that there was no bit set.
  67. size_t find_first_set() const {
  68. #if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_ARM64))
  69. unsigned long result;
  70. bool has_bits_set = (0 != _BitScanForward64(&result, bitset_));
  71. if (!has_bits_set) {
  72. return 0;
  73. }
  74. return result + 1;
  75. #elif defined(_MSC_VER) && defined(_M_IX86)
  76. unsigned long result;
  77. if (static_cast<uint32_t>(bitset_) != 0) {
  78. bool has_bits_set =
  79. (0 != _BitScanForward(&result, static_cast<uint32_t>(bitset_)));
  80. if (!has_bits_set) {
  81. return 0;
  82. }
  83. return result + 1;
  84. } else {
  85. bool has_bits_set =
  86. (0 != _BitScanForward(&result, static_cast<uint32_t>(bitset_ >> 32)));
  87. if (!has_bits_set) {
  88. return 32;
  89. }
  90. return result + 33;
  91. }
  92. #else
  93. return __builtin_ffsll(bitset_);
  94. #endif
  95. }
  96. friend bool operator==(bitset lhs, bitset rhs) noexcept {
  97. return lhs.bitset_ == rhs.bitset_;
  98. }
  99. bitset_type bitset_{0};
  100. };
  101. inline bool operator!=(bitset lhs, bitset rhs) noexcept {
  102. return !(lhs == rhs);
  103. }
  104. } // namespace c10::utils