clearing bit utils (facebook#2301)

Summary:
Pull Request resolved: facebook#2301

n_least_significant_bits
n_most_significant_bits

clear_n_least_significant_bits
set_n_least_significant_bits
clear_n_most_significant_bits
set_n_most_significant_bits

Simple utils that correctly handle corner cases, such as shift == 64.
I looked at the assembly a bit, probably that's ok.
For x86 I used bmi2 where was appropriate.

Reviewed By: Gownta

Differential Revision: D63329499

folly/lang/Bits.h

@@ -67,6 +67,10 @@
 #include <folly/lang/CString.h>
 #include <folly/portability/Builtins.h>
 
+#ifdef __BMI2__
+#include <immintrin.h>
+#endif
+
 #if __has_include(<bit>) && (__cplusplus >= 202002L || (defined(__cpp_lib_bit_cast) && __cpp_lib_bit_cast >= 201806L))
 #include <bit>
 #endif
@@ -106,6 +110,11 @@ constexpr std::make_unsigned_t<Dst> bits_to_unsigned(Src const s) {
   static_assert(std::is_unsigned<Dst>::value, "signed type");
   return static_cast<Dst>(to_unsigned(s));
 }
+
+template <typename T>
+inline constexpr bool supported_in_bits_operations_v =
+    std::is_unsigned_v<T> && sizeof(T) <= 8;
+
 } // namespace detail
 
 /// findFirstSet
@@ -223,6 +232,148 @@ inline constexpr T strictPrevPowTwo(T const v) {
   return v > 1 ? prevPowTwo(T(v - 1)) : T(0);
 }
 
+/// n_least_significant_bits
+/// n_least_significant_bits_fn
+///
+/// Returns an unsigned integer of type T, where n
+/// least significant (right) bits are set and others are not.
+template <class T>
+struct n_least_significant_bits_fn {
+  static_assert(detail::supported_in_bits_operations_v<T>, "");
+
+  FOLLY_NODISCARD constexpr T operator()(std::uint32_t n) const {
+    if (!folly::is_constant_evaluated_or(true)) {
+      compiler_may_unsafely_assume(n <= sizeof(T) * 8);
+
+#ifdef __BMI2__
+      if constexpr (sizeof(T) <= 4) {
+        return static_cast<T>(_bzhi_u32(static_cast<std::uint32_t>(-1), n));
+      }
+      return static_cast<T>(_bzhi_u64(static_cast<std::uint64_t>(-1), n));
+#endif
+    }
+
+    if (sizeof(T) == 8 && n == 64) {
+      return static_cast<T>(-1);
+    }
+    return static_cast<T>((std::uint64_t{1} << n) - 1);
+  }
+};
+
+template <class T>
+inline constexpr n_least_significant_bits_fn<T> n_least_significant_bits;
+
+/// n_most_significant_bits
+/// n_most_significant_bits_fn
+///
+/// Returns an unsigned integer of type T, where n
+/// most significant bits (left) are set and others are not.
+template <class T>
+struct n_most_significant_bits_fn {
+  static_assert(detail::supported_in_bits_operations_v<T>, "");
+
+  FOLLY_NODISCARD constexpr T operator()(std::uint32_t n) const {
+    if (!folly::is_constant_evaluated_or(true)) {
+      compiler_may_unsafely_assume(n <= sizeof(T) * 8);
+
+#ifdef __BMI2__
+      // assembler looks smaller here, if we use bzhi from `set_lowest_n_bits`
+      if constexpr (sizeof(T) == 8) {
+        return static_cast<T>(~n_least_significant_bits<T>(64 - n));
+      }
+#endif
+    }
+
+    if (sizeof(T) == 8 && n == 0) {
+      return 0;
+    }
+    n = sizeof(T) * 8 - n;
+
+    std::uint64_t ones = static_cast<T>(~0);
+    return static_cast<T>(ones << n);
+  }
+};
+
+template <class T>
+inline constexpr n_most_significant_bits_fn<T> n_most_significant_bits;
+
+/// clear_n_least_significant_bits
+/// clear_n_least_significant_bits_fn
+///
+/// Clears n least significant (right) bits. Other bits stay the same.
+struct clear_n_least_significant_bits_fn {
+  template <typename T>
+  FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
+    static_assert(detail::supported_in_bits_operations_v<T>, "");
+
+    // alternative is to do two shifts but that has
+    // a dependency between them, so is likely worse
+    return x & n_most_significant_bits<T>(sizeof(T) * 8 - n);
+  }
+};
+
+inline constexpr clear_n_least_significant_bits_fn
+    clear_n_least_significant_bits;
+
+/// set_n_least_significant_bits
+/// set_n_least_significant_bits_fn
+///
+/// Sets n least significant (right) bits. Other bits stay the same.
+struct set_n_least_significant_bits_fn {
+  template <typename T>
+  FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
+    static_assert(detail::supported_in_bits_operations_v<T>, "");
+
+    // alternative is to do two shifts but that has
+    // a dependency between them, so is likely worse
+    return x | n_least_significant_bits<T>(n);
+  }
+};
+
+inline constexpr set_n_least_significant_bits_fn set_n_least_significant_bits;
+
+/// clear_n_most_significant_bits
+/// clear_n_most_significant_bits_fn
+///
+/// Clears n most significant (left) bits. Other bits stay the same.
+struct clear_n_most_significant_bits_fn {
+  template <typename T>
+  FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
+    static_assert(detail::supported_in_bits_operations_v<T>, "");
+
+    if (!folly::is_constant_evaluated_or(true)) {
+      compiler_may_unsafely_assume(n <= sizeof(T) * 8);
+
+#ifdef __BMI2__
+      if constexpr (sizeof(T) <= 4) {
+        return static_cast<T>(_bzhi_u32(x, sizeof(T) * 8 - n));
+      }
+      return static_cast<T>(_bzhi_u64(x, sizeof(T) * 8 - n));
+#endif
+    }
+
+    // alternative is to do two shifts but that has
+    // a dependency between them, so is likely worse
+    return x & n_least_significant_bits<T>(sizeof(T) * 8 - n);
+  }
+};
+
+inline constexpr clear_n_most_significant_bits_fn clear_n_most_significant_bits;
+
+/// set_n_most_significant_bits
+/// set_n_most_significant_bits_fn
+///
+/// Sets n most significant (left) bits. Other bits stay the same.
+struct set_n_most_significant_bits_fn {
+  template <typename T>
+  FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
+    static_assert(detail::supported_in_bits_operations_v<T>, "");
+    return x | n_most_significant_bits<T>(n);
+  }
+};
+
+inline constexpr set_n_most_significant_bits_fn set_n_most_significant_bits;
+
 /**
  * Endianness detection and manipulation primitives.
  */

folly/lang/test/BitsTest.cpp

@@ -74,6 +74,14 @@ void testEFS() {
   }
 }
 
+template <typename T>
+struct BitsAllUintsTest : ::testing::Test {};
+
+using UintsToTest =
+    ::testing::Types<std::uint8_t, std::uint16_t, std::uint32_t, std::uint64_t>;
+
+TYPED_TEST_SUITE(BitsAllUintsTest, UintsToTest);
+
 } // namespace
 
 TEST(Bits, FindFirstSet) {
@@ -350,4 +358,208 @@ TEST(Bits, LoadUnalignedUB) {
   EXPECT_EQ(0, x);
 }
 
+TYPED_TEST(BitsAllUintsTest, NLeastSignificantBits) {
+  using T = TypeParam;
+
+  static_assert(n_least_significant_bits<T>(0) == 0b0, "");
+  static_assert(n_least_significant_bits<T>(1) == 0b1, "");
+  static_assert(n_least_significant_bits<T>(2) == 0b11, "");
+  static_assert(n_least_significant_bits<T>(3) == 0b111, "");
+  static_assert(n_least_significant_bits<T>(4) == 0b1111, "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= std::min(sizeof(T) * 8, 63UL); ++i) {
+      std::uint64_t expected = (std::uint64_t{1} << i) - 1;
+      T actual = n_least_significant_bits<T>(i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countr_one(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countr_one(expected)) << i;
+        return false;
+      }
+    }
+
+    if (sizeof(T) == 8) {
+      std::uint64_t expected = std::numeric_limits<std::uint64_t>::max();
+      T actual = n_least_significant_bits<T>(64);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << 64;
+        return false;
+      }
+    }
+
+    return true;
+  };
+
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
+TYPED_TEST(BitsAllUintsTest, NMostSignificantBits) {
+  using T = TypeParam;
+
+  constexpr std::size_t kBitSize = sizeof(T) * 8;
+
+  static_assert(
+      n_most_significant_bits<T>(kBitSize) == static_cast<T>(~0b0), "");
+  static_assert(
+      n_most_significant_bits<T>(kBitSize - 1) == static_cast<T>(~0b1), "");
+  static_assert(
+      n_most_significant_bits<T>(kBitSize - 2) == static_cast<T>(~0b11), "");
+  static_assert(
+      n_most_significant_bits<T>(kBitSize - 3) == static_cast<T>(~0b111), "");
+  static_assert(
+      n_most_significant_bits<T>(kBitSize - 4) == static_cast<T>(~0b1111), "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= kBitSize; ++i) {
+      T expected = ~n_least_significant_bits<T>(kBitSize - i);
+      T actual = n_most_significant_bits<T>(i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countl_one(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countl_one(expected)) << i;
+        return false;
+      }
+    }
+    return true;
+  };
+
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
+TYPED_TEST(BitsAllUintsTest, ClearNLeastSignificantBits) {
+  using T = TypeParam;
+
+  constexpr std::size_t kBitSize = sizeof(T) * 8;
+
+  static_assert(clear_n_least_significant_bits(T{0b11U}, 1U) == 0b10U, "");
+  static_assert(clear_n_least_significant_bits(T{0b101U}, 1U) == 0b100U, "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= kBitSize; ++i) {
+      T expected = n_most_significant_bits<T>(kBitSize - i);
+      T actual = clear_n_least_significant_bits(static_cast<T>(-1), i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countr_zero(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countr_zero(expected)) << i;
+        return false;
+      }
+    }
+    return true;
+  };
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
+TYPED_TEST(BitsAllUintsTest, SetNLeastSignificantBits) {
+  using T = TypeParam;
+
+  constexpr std::size_t kBitSize = sizeof(T) * 8;
+
+  static_assert(set_n_least_significant_bits(T{0b10U}, 1U) == 0b11U, "");
+  static_assert(set_n_least_significant_bits(T{0b100U}, 1U) == 0b101U, "");
+  static_assert(set_n_least_significant_bits(T{0b100U}, 2U) == 0b111U, "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= kBitSize; ++i) {
+      T expected = n_least_significant_bits<T>(i);
+      T actual = set_n_least_significant_bits(T{}, i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countr_one(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countr_one(expected)) << i;
+        return false;
+      }
+    }
+    return true;
+  };
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
+TYPED_TEST(BitsAllUintsTest, ClearNMostSignificantBits) {
+  using T = TypeParam;
+
+  constexpr std::size_t kBitSize = sizeof(T) * 8;
+
+  static_assert(
+      clear_n_most_significant_bits(T{0b101U}, kBitSize - 1) == 0b1U, "");
+  static_assert(
+      clear_n_most_significant_bits(T{0b1100U}, kBitSize - 3) == 0b100U, "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= kBitSize; ++i) {
+      T expected = n_least_significant_bits<T>(kBitSize - i);
+      T actual = clear_n_most_significant_bits(static_cast<T>(-1), i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countl_zero(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countl_zero(expected)) << i;
+        return false;
+      }
+    }
+    return true;
+  };
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
+TYPED_TEST(BitsAllUintsTest, SetNMostSignificantBits) {
+  using T = TypeParam;
+
+  constexpr std::size_t kBitSize = sizeof(T) * 8;
+
+  static_assert(
+      set_n_most_significant_bits(T{0b1}, kBitSize - 2) ==
+          static_cast<T>(~0b10),
+      "");
+  static_assert(
+      set_n_most_significant_bits(T{0b1100U}, kBitSize - 3) ==
+          static_cast<T>(~0b11),
+      "");
+
+  auto test = [] {
+    for (std::uint32_t i = 0; i <= kBitSize; ++i) {
+      T expected = n_most_significant_bits<T>(i);
+      T actual = set_n_most_significant_bits(static_cast<T>(0), i);
+      if (expected != actual) {
+        EXPECT_EQ(expected, actual) << i;
+        return false;
+      }
+      if (std::countl_one(expected) != static_cast<int>(i)) {
+        EXPECT_EQ(i, std::countl_one(expected)) << i;
+        return false;
+      }
+    }
+    return true;
+  };
+  static_assert(test(), "");
+
+  // runtime can use a different implementation
+  EXPECT_TRUE(test());
+}
+
 } // namespace folly