diff --git a/folly/lang/Bits.h b/folly/lang/Bits.h index 981feab9be7..2380d516dd3 100644 --- a/folly/lang/Bits.h +++ b/folly/lang/Bits.h @@ -67,6 +67,10 @@ #include #include +#ifdef __BMI2__ +#include +#endif + #if __has_include() && (__cplusplus >= 202002L || (defined(__cpp_lib_bit_cast) && __cpp_lib_bit_cast >= 201806L)) #include #endif @@ -106,6 +110,11 @@ constexpr std::make_unsigned_t bits_to_unsigned(Src const s) { static_assert(std::is_unsigned::value, "signed type"); return static_cast(to_unsigned(s)); } + +template +inline constexpr bool supported_in_bits_operations_v = + std::is_unsigned_v && 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 +struct n_least_significant_bits_fn { + static_assert(detail::supported_in_bits_operations_v, ""); + + 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(_bzhi_u32(static_cast(-1), n)); + } + return static_cast(_bzhi_u64(static_cast(-1), n)); +#endif + } + + if (sizeof(T) == 8 && n == 64) { + return static_cast(-1); + } + return static_cast((std::uint64_t{1} << n) - 1); + } +}; + +template +inline constexpr n_least_significant_bits_fn 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 +struct n_most_significant_bits_fn { + static_assert(detail::supported_in_bits_operations_v, ""); + + 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(~n_least_significant_bits(64 - n)); + } +#endif + } + + if (sizeof(T) == 8 && n == 0) { + return 0; + } + n = sizeof(T) * 8 - n; + + std::uint64_t ones = static_cast(~0); + return static_cast(ones << n); + } +}; + +template +inline constexpr n_most_significant_bits_fn 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 + FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const { + static_assert(detail::supported_in_bits_operations_v, ""); + + // alternative is to do two shifts but that has + // a dependency between them, so is likely worse + return x & n_most_significant_bits(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 + FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const { + static_assert(detail::supported_in_bits_operations_v, ""); + + // alternative is to do two shifts but that has + // a dependency between them, so is likely worse + return x | n_least_significant_bits(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 + FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const { + static_assert(detail::supported_in_bits_operations_v, ""); + + 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(_bzhi_u32(x, sizeof(T) * 8 - n)); + } + return static_cast(_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(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 + FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const { + static_assert(detail::supported_in_bits_operations_v, ""); + return x | n_most_significant_bits(n); + } +}; + +inline constexpr set_n_most_significant_bits_fn set_n_most_significant_bits; + /** * Endianness detection and manipulation primitives. */ diff --git a/folly/lang/test/BitsTest.cpp b/folly/lang/test/BitsTest.cpp index 5888c748872..350fb114d23 100644 --- a/folly/lang/test/BitsTest.cpp +++ b/folly/lang/test/BitsTest.cpp @@ -74,6 +74,14 @@ void testEFS() { } } +template +struct BitsAllUintsTest : ::testing::Test {}; + +using UintsToTest = + ::testing::Types; + +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(0) == 0b0, ""); + static_assert(n_least_significant_bits(1) == 0b1, ""); + static_assert(n_least_significant_bits(2) == 0b11, ""); + static_assert(n_least_significant_bits(3) == 0b111, ""); + static_assert(n_least_significant_bits(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(i); + if (expected != actual) { + EXPECT_EQ(expected, actual) << i; + return false; + } + if (std::countr_one(expected) != static_cast(i)) { + EXPECT_EQ(i, std::countr_one(expected)) << i; + return false; + } + } + + if (sizeof(T) == 8) { + std::uint64_t expected = std::numeric_limits::max(); + T actual = n_least_significant_bits(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(kBitSize) == static_cast(~0b0), ""); + static_assert( + n_most_significant_bits(kBitSize - 1) == static_cast(~0b1), ""); + static_assert( + n_most_significant_bits(kBitSize - 2) == static_cast(~0b11), ""); + static_assert( + n_most_significant_bits(kBitSize - 3) == static_cast(~0b111), ""); + static_assert( + n_most_significant_bits(kBitSize - 4) == static_cast(~0b1111), ""); + + auto test = [] { + for (std::uint32_t i = 0; i <= kBitSize; ++i) { + T expected = ~n_least_significant_bits(kBitSize - i); + T actual = n_most_significant_bits(i); + if (expected != actual) { + EXPECT_EQ(expected, actual) << i; + return false; + } + if (std::countl_one(expected) != static_cast(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(kBitSize - i); + T actual = clear_n_least_significant_bits(static_cast(-1), i); + if (expected != actual) { + EXPECT_EQ(expected, actual) << i; + return false; + } + if (std::countr_zero(expected) != static_cast(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(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(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(kBitSize - i); + T actual = clear_n_most_significant_bits(static_cast(-1), i); + if (expected != actual) { + EXPECT_EQ(expected, actual) << i; + return false; + } + if (std::countl_zero(expected) != static_cast(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(~0b10), + ""); + static_assert( + set_n_most_significant_bits(T{0b1100U}, kBitSize - 3) == + static_cast(~0b11), + ""); + + auto test = [] { + for (std::uint32_t i = 0; i <= kBitSize; ++i) { + T expected = n_most_significant_bits(i); + T actual = set_n_most_significant_bits(static_cast(0), i); + if (expected != actual) { + EXPECT_EQ(expected, actual) << i; + return false; + } + if (std::countl_one(expected) != static_cast(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