1 files changed, 109 insertions, 145 deletions

diff --git a/src/__support/math_extras.h b/src/__support/math_extras.h
index 8ec30396ffdb..4bd871957406 100644
--- a/src/__support/math_extras.h
+++ b/src/__support/math_extras.h
@@ -10,187 +10,151 @@
 #ifndef LLVM_LIBC_SRC___SUPPORT_MATH_EXTRAS_H
 #define LLVM_LIBC_SRC___SUPPORT_MATH_EXTRAS_H
 
-#include "src/__support/CPP/limits.h"        // CHAR_BIT
-#include "src/__support/CPP/type_traits.h"   // is_unsigned_v
+#include "src/__support/CPP/bit.h"         // countl_one, countr_zero
+#include "src/__support/CPP/limits.h"      // CHAR_BIT, numeric_limits
+#include "src/__support/CPP/type_traits.h" // is_unsigned_v, is_constant_evaluated
 #include "src/__support/macros/attributes.h" // LIBC_INLINE
-#include "src/__support/macros/config.h"     // LIBC_HAS_BUILTIN
 
 namespace LIBC_NAMESPACE {
 
 // Create a bitmask with the count right-most bits set to 1, and all other bits
 // set to 0.  Only unsigned types are allowed.
 template <typename T, size_t count>
-LIBC_INLINE constexpr T mask_trailing_ones() {
-  static_assert(cpp::is_unsigned_v<T>);
-  constexpr unsigned t_bits = CHAR_BIT * sizeof(T);
-  static_assert(count <= t_bits && "Invalid bit index");
-  // It's important not to initialize T with -1, since T may be BigInt which
-  // will take -1 as a uint64_t and only initialize the low 64 bits.
-  constexpr T all_zeroes(0);
-  constexpr T all_ones(~all_zeroes); // bitwise NOT performs integer promotion.
-  return count == 0 ? 0 : (all_ones >> (t_bits - count));
+LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+mask_trailing_ones() {
+  constexpr unsigned T_BITS = CHAR_BIT * sizeof(T);
+  static_assert(count <= T_BITS && "Invalid bit index");
+  return count == 0 ? 0 : (T(-1) >> (T_BITS - count));
 }
 
 // Create a bitmask with the count left-most bits set to 1, and all other bits
 // set to 0.  Only unsigned types are allowed.
 template <typename T, size_t count>
-LIBC_INLINE constexpr T mask_leading_ones() {
-  constexpr T mask(mask_trailing_ones<T, CHAR_BIT * sizeof(T) - count>());
-  return T(~mask); // bitwise NOT performs integer promotion.
+LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+mask_leading_ones() {
+  return T(~mask_trailing_ones<T, CHAR_BIT * sizeof(T) - count>());
 }
 
-// Add with carry
-template <typename T> struct SumCarry {
-  T sum;
-  T carry;
-};
-
-// This version is always valid for constexpr.
-template <typename T>
-LIBC_INLINE constexpr cpp::enable_if_t<
-    cpp::is_integral_v<T> && cpp::is_unsigned_v<T>, SumCarry<T>>
-add_with_carry_const(T a, T b, T carry_in) {
-  T tmp = a + carry_in;
-  T sum = b + tmp;
-  T carry_out = (sum < b) + (tmp < a);
-  return {sum, carry_out};
-}
-
-// This version is not always valid for constepxr because it's overriden below
-// if builtins are available.
-template <typename T>
-LIBC_INLINE cpp::enable_if_t<cpp::is_integral_v<T> && cpp::is_unsigned_v<T>,
-                             SumCarry<T>>
-add_with_carry(T a, T b, T carry_in) {
-  return add_with_carry_const<T>(a, b, carry_in);
-}
-
-#if LIBC_HAS_BUILTIN(__builtin_addc)
-// https://clang.llvm.org/docs/LanguageExtensions.html#multiprecision-arithmetic-builtins
-
-template <>
-LIBC_INLINE SumCarry<unsigned char>
-add_with_carry<unsigned char>(unsigned char a, unsigned char b,
-                              unsigned char carry_in) {
-  SumCarry<unsigned char> result{0, 0};
-  result.sum = __builtin_addcb(a, b, carry_in, &result.carry);
-  return result;
-}
-
-template <>
-LIBC_INLINE SumCarry<unsigned short>
-add_with_carry<unsigned short>(unsigned short a, unsigned short b,
-                               unsigned short carry_in) {
-  SumCarry<unsigned short> result{0, 0};
-  result.sum = __builtin_addcs(a, b, carry_in, &result.carry);
-  return result;
+// Create a bitmask with the count right-most bits set to 0, and all other bits
+// set to 1.  Only unsigned types are allowed.
+template <typename T, size_t count>
+LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+mask_trailing_zeros() {
+  return mask_leading_ones<T, CHAR_BIT * sizeof(T) - count>();
 }
 
-template <>
-LIBC_INLINE SumCarry<unsigned int>
-add_with_carry<unsigned int>(unsigned int a, unsigned int b,
-                             unsigned int carry_in) {
-  SumCarry<unsigned int> result{0, 0};
-  result.sum = __builtin_addc(a, b, carry_in, &result.carry);
-  return result;
+// Create a bitmask with the count left-most bits set to 0, and all other bits
+// set to 1.  Only unsigned types are allowed.
+template <typename T, size_t count>
+LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+mask_leading_zeros() {
+  return mask_trailing_ones<T, CHAR_BIT * sizeof(T) - count>();
 }
 
-template <>
-LIBC_INLINE SumCarry<unsigned long>
-add_with_carry<unsigned long>(unsigned long a, unsigned long b,
-                              unsigned long carry_in) {
-  SumCarry<unsigned long> result{0, 0};
-  result.sum = __builtin_addcl(a, b, carry_in, &result.carry);
-  return result;
+// Returns whether 'a + b' overflows, the result is stored in 'res'.
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr bool add_overflow(T a, T b, T &res) {
+  return __builtin_add_overflow(a, b, &res);
 }
 
-template <>
-LIBC_INLINE SumCarry<unsigned long long>
-add_with_carry<unsigned long long>(unsigned long long a, unsigned long long b,
-                                   unsigned long long carry_in) {
-  SumCarry<unsigned long long> result{0, 0};
-  result.sum = __builtin_addcll(a, b, carry_in, &result.carry);
-  return result;
+// Returns whether 'a - b' overflows, the result is stored in 'res'.
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr bool sub_overflow(T a, T b, T &res) {
+  return __builtin_sub_overflow(a, b, &res);
 }
 
-#endif // LIBC_HAS_BUILTIN(__builtin_addc)
-
-// Subtract with borrow
-template <typename T> struct DiffBorrow {
-  T diff;
-  T borrow;
-};
+#define RETURN_IF(TYPE, BUILTIN)                                               \
+  if constexpr (cpp::is_same_v<T, TYPE>)                                       \
+    return BUILTIN(a, b, carry_in, carry_out);
 
-// This version is always valid for constexpr.
+// Returns the result of 'a + b' taking into account 'carry_in'.
+// The carry out is stored in 'carry_out' it not 'nullptr', dropped otherwise.
+// We keep the pass by pointer interface for consistency with the intrinsic.
 template <typename T>
-LIBC_INLINE constexpr cpp::enable_if_t<
-    cpp::is_integral_v<T> && cpp::is_unsigned_v<T>, DiffBorrow<T>>
-sub_with_borrow_const(T a, T b, T borrow_in) {
-  T tmp = a - b;
-  T diff = tmp - borrow_in;
-  T borrow_out = (diff > tmp) + (tmp > a);
-  return {diff, borrow_out};
-}
-
-// This version is not always valid for constepxr because it's overriden below
-// if builtins are available.
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+add_with_carry(T a, T b, T carry_in, T &carry_out) {
+  if constexpr (!cpp::is_constant_evaluated()) {
+#if __has_builtin(__builtin_addcb)
+    RETURN_IF(unsigned char, __builtin_addcb)
+#elif __has_builtin(__builtin_addcs)
+    RETURN_IF(unsigned short, __builtin_addcs)
+#elif __has_builtin(__builtin_addc)
+    RETURN_IF(unsigned int, __builtin_addc)
+#elif __has_builtin(__builtin_addcl)
+    RETURN_IF(unsigned long, __builtin_addcl)
+#elif __has_builtin(__builtin_addcll)
+    RETURN_IF(unsigned long long, __builtin_addcll)
+#endif
+  }
+  T sum = {};
+  T carry1 = add_overflow(a, b, sum);
+  T carry2 = add_overflow(sum, carry_in, sum);
+  carry_out = carry1 | carry2;
+  return sum;
+}
+
+// Returns the result of 'a - b' taking into account 'carry_in'.
+// The carry out is stored in 'carry_out' it not 'nullptr', dropped otherwise.
+// We keep the pass by pointer interface for consistency with the intrinsic.
 template <typename T>
-LIBC_INLINE cpp::enable_if_t<cpp::is_integral_v<T> && cpp::is_unsigned_v<T>,
-                             DiffBorrow<T>>
-sub_with_borrow(T a, T b, T borrow_in) {
-  return sub_with_borrow_const<T>(a, b, borrow_in);
-}
-
-#if LIBC_HAS_BUILTIN(__builtin_subc)
-// https://clang.llvm.org/docs/LanguageExtensions.html#multiprecision-arithmetic-builtins
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, T>
+sub_with_borrow(T a, T b, T carry_in, T &carry_out) {
+  if constexpr (!cpp::is_constant_evaluated()) {
+#if __has_builtin(__builtin_subcb)
+    RETURN_IF(unsigned char, __builtin_subcb)
+#elif __has_builtin(__builtin_subcs)
+    RETURN_IF(unsigned short, __builtin_subcs)
+#elif __has_builtin(__builtin_subc)
+    RETURN_IF(unsigned int, __builtin_subc)
+#elif __has_builtin(__builtin_subcl)
+    RETURN_IF(unsigned long, __builtin_subcl)
+#elif __has_builtin(__builtin_subcll)
+    RETURN_IF(unsigned long long, __builtin_subcll)
+#endif
+  }
+  T sub = {};
+  T carry1 = sub_overflow(a, b, sub);
+  T carry2 = sub_overflow(sub, carry_in, sub);
+  carry_out = carry1 | carry2;
+  return sub;
+}
+
+#undef RETURN_IF
 
-template <>
-LIBC_INLINE DiffBorrow<unsigned char>
-sub_with_borrow<unsigned char>(unsigned char a, unsigned char b,
-                               unsigned char borrow_in) {
-  DiffBorrow<unsigned char> result{0, 0};
-  result.diff = __builtin_subcb(a, b, borrow_in, &result.borrow);
-  return result;
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, int>
+first_leading_zero(T value) {
+  return value == cpp::numeric_limits<T>::max() ? 0
+                                                : cpp::countl_one(value) + 1;
 }
 
-template <>
-LIBC_INLINE DiffBorrow<unsigned short>
-sub_with_borrow<unsigned short>(unsigned short a, unsigned short b,
-                                unsigned short borrow_in) {
-  DiffBorrow<unsigned short> result{0, 0};
-  result.diff = __builtin_subcs(a, b, borrow_in, &result.borrow);
-  return result;
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, int>
+first_leading_one(T value) {
+  return first_leading_zero(static_cast<T>(~value));
 }
 
-template <>
-LIBC_INLINE DiffBorrow<unsigned int>
-sub_with_borrow<unsigned int>(unsigned int a, unsigned int b,
-                              unsigned int borrow_in) {
-  DiffBorrow<unsigned int> result{0, 0};
-  result.diff = __builtin_subc(a, b, borrow_in, &result.borrow);
-  return result;
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, int>
+first_trailing_zero(T value) {
+  return value == cpp::numeric_limits<T>::max()
+             ? 0
+             : cpp::countr_zero(static_cast<T>(~value)) + 1;
 }
 
-template <>
-LIBC_INLINE DiffBorrow<unsigned long>
-sub_with_borrow<unsigned long>(unsigned long a, unsigned long b,
-                               unsigned long borrow_in) {
-  DiffBorrow<unsigned long> result{0, 0};
-  result.diff = __builtin_subcl(a, b, borrow_in, &result.borrow);
-  return result;
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, int>
+first_trailing_one(T value) {
+  return value == cpp::numeric_limits<T>::max() ? 0
+                                                : cpp::countr_zero(value) + 1;
 }
 
-template <>
-LIBC_INLINE DiffBorrow<unsigned long long>
-sub_with_borrow<unsigned long long>(unsigned long long a, unsigned long long b,
-                                    unsigned long long borrow_in) {
-  DiffBorrow<unsigned long long> result{0, 0};
-  result.diff = __builtin_subcll(a, b, borrow_in, &result.borrow);
-  return result;
+template <typename T>
+[[nodiscard]] LIBC_INLINE constexpr cpp::enable_if_t<cpp::is_unsigned_v<T>, int>
+count_zeros(T value) {
+  return cpp::popcount<T>(static_cast<T>(~value));
 }
 
-#endif // LIBC_HAS_BUILTIN(__builtin_subc)
-
 } // namespace LIBC_NAMESPACE
 
 #endif // LLVM_LIBC_SRC___SUPPORT_MATH_EXTRAS_H