[libc][__support][bit] add count_zeros

[nickdesaulniers]

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

[llvmbot]

@llvm/pr-subscribers-libc

Author: Nick Desaulniers (nickdesaulniers)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/82076.diff

2 Files Affected:

• (modified) libc/src/__support/CPP/bit.h (+28)
• (modified) libc/test/src/__support/CPP/bit_test.cpp (+13)

diff --git a/libc/src/__support/CPP/bit.h b/libc/src/__support/CPP/bit.h
index f5e50262371f26..73c5949d86423b 100644
--- a/libc/src/__support/CPP/bit.h
+++ b/libc/src/__support/CPP/bit.h
@@ -248,6 +248,34 @@ 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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.