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

[Nick Desaulniers via libc-commits]

https://github.com/nickdesaulniers created https://github.com/llvm/llvm-project/pull/82076

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


>From 2d4f7c7b8524b580c0100e5e52e7eae6b1408b18 Mon Sep 17 00:00:00 2001
From: Nick Desaulniers <ndesaulniers at google.com>
Date: Fri, 16 Feb 2024 16:15:55 -0800
Subject: [PATCH] [libc][__support][bit] add count_zeros

Will be useful for implementing C23 stdbit.h's stdc_count_zeros and
stdc_count_ones.
---
 libc/src/__support/CPP/bit.h             | 28 ++++++++++++++++++++++++
 libc/test/src/__support/CPP/bit_test.cpp | 13 +++++++++++
 2 files changed, 41 insertions(+)

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