gh-130317: Fix PyFloat_Pack/Unpack[24] for NaN's with payload (#130452)

Co-authored-by: Victor Stinner <vstinner@python.org>

5 files changed, 120 insertions, 8 deletions

diff --git a/Lib/test/test_capi/test_float.py b/Lib/test/test_capi/test_float.py
index 92c9877..c857959 100644
--- a/Lib/test/test_capi/test_float.py
+++ b/Lib/test/test_capi/test_float.py
@@ -1,4 +1,5 @@
 import math
+import random
 import sys
 import unittest
 import warnings
@@ -178,6 +179,39 @@ class CAPIFloatTest(unittest.TestCase):
                         else:
                             self.assertEqual(value2, value)
 
+    @unittest.skipUnless(HAVE_IEEE_754, "requires IEEE 754")
+    def test_pack_unpack_roundtrip_for_nans(self):
+        pack = _testcapi.float_pack
+        unpack = _testcapi.float_unpack
+        for _ in range(1000):
+            for size in (2, 4, 8):
+                sign = random.randint(0, 1)
+                signaling = random.randint(0, 1)
+                quiet = int(not signaling)
+                if size == 8:
+                    payload = random.randint(signaling, 1 << 50)
+                    i = (sign << 63) + (0x7ff << 52) + (quiet << 51) + payload
+                elif size == 4:
+                    payload = random.randint(signaling, 1 << 21)
+                    i = (sign << 31) + (0xff << 23) + (quiet << 22) + payload
+                elif size == 2:
+                    payload = random.randint(signaling, 1 << 8)
+                    i = (sign << 15) + (0x1f << 10) + (quiet << 9) + payload
+                data = bytes.fromhex(f'{i:x}')
+                for endian in (BIG_ENDIAN, LITTLE_ENDIAN):
+                    with self.subTest(data=data, size=size, endian=endian):
+                        data1 = data if endian == BIG_ENDIAN else data[::-1]
+                        value = unpack(data1, endian)
+                        if signaling and sys.platform == 'win32':
+                            # On this platform sNaN becomes qNaN when returned
+                            # from function.  That's a known bug, e.g.
+                            # https://developercommunity.visualstudio.com/t/155064
+                            # (see also gh-130317).
+                            value = _testcapi.float_set_snan(value)
+                        data2 = pack(size, value, endian)
+                        self.assertTrue(math.isnan(value))
+                        self.assertEqual(data1, data2)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2025-02-22-13-07-06.gh-issue-130317.tnxd0I.rst b/Misc/NEWS.d/next/Library/2025-02-22-13-07-06.gh-issue-130317.tnxd0I.rst
new file mode 100644
index 0000000..ab69f88
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2025-02-22-13-07-06.gh-issue-130317.tnxd0I.rst
@@ -0,0 +1,4 @@
+Fix :c:func:`PyFloat_Pack2` and :c:func:`PyFloat_Unpack2` for NaN's with
+payload.  This corrects round-trip for :func:`struct.unpack` and
+:func:`struct.pack` in case of the IEEE 754 binary16 "half precision" type.
+Patch by Sergey B Kirpichev.
diff --git a/Modules/_testcapi/clinic/float.c.h b/Modules/_testcapi/clinic/float.c.h
index d5a00c8..0710e4d 100644
--- a/Modules/_testcapi/clinic/float.c.h
+++ b/Modules/_testcapi/clinic/float.c.h
@@ -81,4 +81,13 @@ _testcapi_float_unpack(PyObject *module, PyObject *const *args, Py_ssize_t nargs
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=b43dfd3a77fe04ba input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_testcapi_float_set_snan__doc__,
+"float_set_snan($module, obj, /)\n"
+"--\n"
+"\n"
+"Make a signaling NaN.");
+
+#define _TESTCAPI_FLOAT_SET_SNAN_METHODDEF    \
+    {"float_set_snan", (PyCFunction)_testcapi_float_set_snan, METH_O, _testcapi_float_set_snan__doc__},
+/*[clinic end generated code: output=1b0e9b05e1f50712 input=a9049054013a1b77]*/
diff --git a/Modules/_testcapi/float.c b/Modules/_testcapi/float.c
index e386913..2feeb20 100644
--- a/Modules/_testcapi/float.c
+++ b/Modules/_testcapi/float.c
@@ -157,9 +157,39 @@ test_string_to_double(PyObject *self, PyObject *Py_UNUSED(ignored))
 }
 
 
+/*[clinic input]
+_testcapi.float_set_snan
+
+    obj: object
+    /
+
+Make a signaling NaN.
+[clinic start generated code]*/
+
+static PyObject *
+_testcapi_float_set_snan(PyObject *module, PyObject *obj)
+/*[clinic end generated code: output=f43778a70f60aa4b input=c1269b0f88ef27ac]*/
+{
+    if (!PyFloat_Check(obj)) {
+        PyErr_SetString(PyExc_ValueError, "float-point number expected");
+        return NULL;
+    }
+    double d = ((PyFloatObject *)obj)->ob_fval;
+    if (!isnan(d)) {
+        PyErr_SetString(PyExc_ValueError, "nan expected");
+        return NULL;
+    }
+    uint64_t v;
+    memcpy(&v, &d, 8);
+    v &= ~(1ULL << 51); /* make sNaN */
+    memcpy(&d, &v, 8);
+    return PyFloat_FromDouble(d);
+}
+
 static PyMethodDef test_methods[] = {
     _TESTCAPI_FLOAT_PACK_METHODDEF
     _TESTCAPI_FLOAT_UNPACK_METHODDEF
+    _TESTCAPI_FLOAT_SET_SNAN_METHODDEF
     {"test_string_to_double", test_string_to_double, METH_NOARGS},
     {NULL},
 };
diff --git a/Objects/floatobject.c b/Objects/floatobject.c
index 87a00bf..e0a8f0c 100644
--- a/Objects/floatobject.c
+++ b/Objects/floatobject.c
@@ -2021,14 +2021,13 @@ PyFloat_Pack2(double x, char *data, int le)
         bits = 0;
     }
     else if (isnan(x)) {
-        /* There are 2046 distinct half-precision NaNs (1022 signaling and
-           1024 quiet), but there are only two quiet NaNs that don't arise by
-           quieting a signaling NaN; we get those by setting the topmost bit
-           of the fraction field and clearing all other fraction bits. We
-           choose the one with the appropriate sign. */
         sign = (copysign(1.0, x) == -1.0);
         e = 0x1f;
-        bits = 512;
+
+        uint64_t v;
+        memcpy(&v, &x, sizeof(v));
+        v &= 0xffc0000000000ULL;
+        bits = (unsigned short)(v >> 42); /* NaN's type & payload */
     }
     else {
         sign = (x < 0.0);
@@ -2192,6 +2191,21 @@ PyFloat_Pack4(double x, char *data, int le)
         if (isinf(y) && !isinf(x))
             goto Overflow;
 
+        /* correct y if x was a sNaN, transformed to qNaN by conversion */
+        if (isnan(x)) {
+            uint64_t v;
+
+            memcpy(&v, &x, 8);
+            if ((v & (1ULL << 51)) == 0) {
+                union float_val {
+                    float f;
+                    uint32_t u32;
+                } *py = (union float_val *)&y;
+
+                py->u32 &= ~(1 << 22); /* make sNaN */
+            }
+        }
+
         unsigned char s[sizeof(float)];
         memcpy(s, &y, sizeof(float));
 
@@ -2374,7 +2388,11 @@ PyFloat_Unpack2(const char *data, int le)
         }
         else {
             /* NaN */
-            return sign ? -fabs(Py_NAN) : fabs(Py_NAN);
+            uint64_t v = sign ? 0xfff0000000000000ULL : 0x7ff0000000000000ULL;
+
+            v += (uint64_t)f << 42; /* add NaN's type & payload */
+            memcpy(&x, &v, sizeof(v));
+            return x;
         }
     }
 
@@ -2470,6 +2488,23 @@ PyFloat_Unpack4(const char *data, int le)
             memcpy(&x, p, 4);
         }
 
+        /* return sNaN double if x was sNaN float */
+        if (isnan(x)) {
+            uint32_t v;
+            memcpy(&v, &x, 4);
+
+            if ((v & (1 << 22)) == 0) {
+                double y = x; /* will make qNaN double */
+                union double_val {
+                    double d;
+                    uint64_t u64;
+                } *py = (union double_val *)&y;
+
+                py->u64 &= ~(1ULL << 51); /* make sNaN */
+                return y;
+            }
+        }
+
         return x;
     }
 }