[libc-commits] [libc] [libc][NFC] Simplify logic in `sqrt` (PR #80426)

[Guillaume Chatelet via libc-commits]

https://github.com/gchatelet updated https://github.com/llvm/llvm-project/pull/80426

>From 95c821c01b08e42c8a61239bac16fd56c598bde0 Mon Sep 17 00:00:00 2001
From: Guillaume Chatelet <gchatelet at google.com>
Date: Fri, 2 Feb 2024 12:43:31 +0000
Subject: [PATCH 1/2] [libc][NFC] Simplify logic in `sqrt`

---
 libc/src/__support/FPUtil/generic/sqrt.h         | 16 +++++-----------
 .../FPUtil/generic/sqrt_80_bit_long_double.h     | 16 +++++-----------
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/libc/src/__support/FPUtil/generic/sqrt.h b/libc/src/__support/FPUtil/generic/sqrt.h
index 702de3f04a9bf..15ac1d3545e95 100644
--- a/libc/src/__support/FPUtil/generic/sqrt.h
+++ b/libc/src/__support/FPUtil/generic/sqrt.h
@@ -78,20 +78,14 @@ LIBC_INLINE cpp::enable_if_t<cpp::is_floating_point_v<T>, T> sqrt(T x) {
 
     FPBits_t bits(x);
 
-    if (bits.is_inf_or_nan()) {
-      if (bits.is_neg() && (bits.get_mantissa() == 0)) {
-        // sqrt(-Inf) = NaN
-        return FLT_NAN;
-      } else {
-        // sqrt(NaN) = NaN
-        // sqrt(+Inf) = +Inf
-        return x;
-      }
-    } else if (bits.is_zero()) {
+    if (bits == FPBits_t::inf(Sign::POS) || bits.is_zero()) {
+      // sqrt(+Inf) = +Inf
       // sqrt(+0) = +0
       // sqrt(-0) = -0
       return x;
-    } else if (bits.is_neg()) {
+    } else if (bits.is_inf_or_nan() || bits.is_neg()) {
+      // sqrt(-Inf) = NaN
+      // sqrt(NaN) = NaN
       // sqrt( negative numbers ) = NaN
       return FLT_NAN;
     } else {
diff --git a/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h b/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
index 74a536c04817c..2146019011986 100644
--- a/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
+++ b/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
@@ -43,20 +43,14 @@ LIBC_INLINE long double sqrt(long double x) {
 
   LDBits bits(x);
 
-  if (bits.is_inf_or_nan()) {
-    if (bits.is_neg() && (bits.get_mantissa() == 0)) {
-      // sqrt(-Inf) = NaN
-      return LDNAN;
-    } else {
-      // sqrt(NaN) = NaN
-      // sqrt(+Inf) = +Inf
-      return x;
-    }
-  } else if (bits.is_zero()) {
+  if (bits == LDBits::inf(Sign::POS) || bits.is_zero()) {
+    // sqrt(+Inf) = +Inf
     // sqrt(+0) = +0
     // sqrt(-0) = -0
     return x;
-  } else if (bits.is_neg()) {
+  } else if (bits.is_inf_or_nan() || bits.is_neg()) {
+    // sqrt(-Inf) = NaN
+    // sqrt(NaN) = NaN
     // sqrt( negative numbers ) = NaN
     return LDNAN;
   } else {

>From 9a5af5f2314f1a050dddb45e242ae05a67192b10 Mon Sep 17 00:00:00 2001
From: Guillaume Chatelet <gchatelet at google.com>
Date: Fri, 2 Feb 2024 15:56:56 +0000
Subject: [PATCH 2/2] Address comments

---
 libc/src/__support/FPUtil/generic/sqrt.h                 | 9 +++++----
 .../__support/FPUtil/generic/sqrt_80_bit_long_double.h   | 9 +++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/libc/src/__support/FPUtil/generic/sqrt.h b/libc/src/__support/FPUtil/generic/sqrt.h
index 15ac1d3545e95..9e8ed305eeca5 100644
--- a/libc/src/__support/FPUtil/generic/sqrt.h
+++ b/libc/src/__support/FPUtil/generic/sqrt.h
@@ -78,15 +78,16 @@ LIBC_INLINE cpp::enable_if_t<cpp::is_floating_point_v<T>, T> sqrt(T x) {
 
     FPBits_t bits(x);
 
-    if (bits == FPBits_t::inf(Sign::POS) || bits.is_zero()) {
+    if (bits == FPBits_t::inf(Sign::POS) || bits.is_zero() || bits.is_nan()) {
       // sqrt(+Inf) = +Inf
       // sqrt(+0) = +0
       // sqrt(-0) = -0
+      // sqrt(NaN) = NaN
+      // sqrt(-NaN) = -NaN
       return x;
-    } else if (bits.is_inf_or_nan() || bits.is_neg()) {
+    } else if (bits.is_neg()) {
       // sqrt(-Inf) = NaN
-      // sqrt(NaN) = NaN
-      // sqrt( negative numbers ) = NaN
+      // sqrt(-x) = NaN
       return FLT_NAN;
     } else {
       int x_exp = bits.get_exponent();
diff --git a/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h b/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
index 2146019011986..4b2af10256176 100644
--- a/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
+++ b/libc/src/__support/FPUtil/generic/sqrt_80_bit_long_double.h
@@ -43,15 +43,16 @@ LIBC_INLINE long double sqrt(long double x) {
 
   LDBits bits(x);
 
-  if (bits == LDBits::inf(Sign::POS) || bits.is_zero()) {
+  if (bits == LDBits::inf(Sign::POS) || bits.is_zero() || bits.is_nan()) {
     // sqrt(+Inf) = +Inf
     // sqrt(+0) = +0
     // sqrt(-0) = -0
+    // sqrt(NaN) = NaN
+    // sqrt(-NaN) = -NaN
     return x;
-  } else if (bits.is_inf_or_nan() || bits.is_neg()) {
+  } else if (bits.is_neg()) {
     // sqrt(-Inf) = NaN
-    // sqrt(NaN) = NaN
-    // sqrt( negative numbers ) = NaN
+    // sqrt(-x) = NaN
     return LDNAN;
   } else {
     int x_exp = bits.get_explicit_exponent();