[libc-commits] [libc] ed4bdb8 - [libc][__support][bit] add count_zeros (#82076)

[via libc-commits]

Author: Nick Desaulniers
Date: 2024-02-20T12:38:23-08:00
New Revision: ed4bdb86b084bf633770136d005426adeeb2cd57

URL: https://github.com/llvm/llvm-project/commit/ed4bdb86b084bf633770136d005426adeeb2cd57
DIFF: https://github.com/llvm/llvm-project/commit/ed4bdb86b084bf633770136d005426adeeb2cd57.diff

LOG: [libc][__support][bit] add count_zeros (#82076)


Will be useful for implementing C23 stdbit.h's stdc_count_zeros and
stdc_count_ones.

Added: 
    

Modified: 
    libc/src/__support/CPP/bit.h
    libc/test/src/__support/CPP/bit_test.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/CPP/bit.h b/libc/src/__support/CPP/bit.h
index f5e50262371f26..7d11e7d5c497e0 100644
--- a/libc/src/__support/CPP/bit.h
+++ b/libc/src/__support/CPP/bit.h
@@ -248,6 +248,35 @@ template <typename T, typename = cpp::enable_if_t<cpp::is_unsigned_v<T>>>
   return value == cpp::numeric_limits<T>::max() ? 0 : countr_zero(value) + 1;
 }
 
+/// Count number of 1's aka population count or hamming weight.
+///
+/// Only unsigned integral types are allowed.
+template <typename T, typename = cpp::enable_if_t<cpp::is_unsigned_v<T>>>
+[[nodiscard]] LIBC_INLINE constexpr int count_ones(T value) {
+  int count = 0;
+  for (int i = 0; i != cpp::numeric_limits<T>::digits; ++i)
+    if ((value >> i) & 0x1)
+      ++count;
+  return count;
+}
+#define ADD_SPECIALIZATION(TYPE, BUILTIN)                                      \
+  template <>                                                                  \
+  [[nodiscard]] LIBC_INLINE constexpr int count_ones<TYPE>(TYPE value) {       \
+    return BUILTIN(value);                                                     \
+  }
+ADD_SPECIALIZATION(unsigned char, __builtin_popcount)
+ADD_SPECIALIZATION(unsigned short, __builtin_popcount)
+ADD_SPECIALIZATION(unsigned, __builtin_popcount)
+ADD_SPECIALIZATION(unsigned long, __builtin_popcountl)
+ADD_SPECIALIZATION(unsigned long long, __builtin_popcountll)
+// TODO: 128b specializations?
+#undef ADD_SPECIALIZATION
+
+template <typename T, typename = cpp::enable_if_t<cpp::is_unsigned_v<T>>>
+[[nodiscard]] LIBC_INLINE constexpr int count_zeros(T value) {
+  return count_ones<T>(static_cast<T>(~value));
+}
+
 } // namespace LIBC_NAMESPACE::cpp
 
 #endif // LLVM_LIBC_SRC___SUPPORT_CPP_BIT_H

diff  --git a/libc/test/src/__support/CPP/bit_test.cpp b/libc/test/src/__support/CPP/bit_test.cpp
index 5d1f451776a5fe..115a5d505c4b7a 100644
--- a/libc/test/src/__support/CPP/bit_test.cpp
+++ b/libc/test/src/__support/CPP/bit_test.cpp
@@ -232,4 +232,17 @@ TYPED_TEST(LlvmLibcBitTest, FirstTrailingOne, UnsignedTypes) {
     EXPECT_EQ(first_trailing_one<T>(T(1) << i), i + 1);
 }
 
+TYPED_TEST(LlvmLibcBitTest, CountZeros, UnsignedTypes) {
+  EXPECT_EQ(count_zeros(T(0)), cpp::numeric_limits<T>::digits);
+  for (int i = 0; i != cpp::numeric_limits<T>::digits; ++i)
+    EXPECT_EQ(count_zeros<T>(cpp::numeric_limits<T>::max() >> i), i);
+}
+
+TYPED_TEST(LlvmLibcBitTest, CountOnes, UnsignedTypes) {
+  EXPECT_EQ(count_ones(T(0)), 0);
+  for (int i = 0; i != cpp::numeric_limits<T>::digits; ++i)
+    EXPECT_EQ(count_ones<T>(cpp::numeric_limits<T>::max() >> i),
+              cpp::numeric_limits<T>::digits - i);
+}
+
 } // namespace LIBC_NAMESPACE::cpp