Address Irit's review && Add guards for NaN

python · Eclips4 · Jan 19, 2025 · Jan 17, 2025 · Jan 17, 2025 · Jan 17, 2025
commit 8363e9d270d95b7d23a0cda5dff6a9448bc28ea5
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -1362,19 +1362,47 @@ def binary_op_add_extend():
         self.assert_specialized(binary_op_add_extend, "BINARY_OP_EXTEND")
         self.assert_no_opcode(binary_op_add_extend, "BINARY_OP")
 
-        def compactlong_lhs(arg):
-            42 / arg
-        def float_lhs(arg):
-            42.0 / arg
-
-        with self.assertRaises(ZeroDivisionError):
-            compactlong_lhs(0)
-        with self.assertRaises(ZeroDivisionError):
-            compactlong_lhs(0.0)
-        with self.assertRaises(ZeroDivisionError):
-            float_lhs(0.0)
-        with self.assertRaises(ZeroDivisionError):
-            float_lhs(0)
+        def binary_op_zero_division():
+            def compactlong_lhs(arg):
+                42 / arg
+            def float_lhs(arg):
+                42.0 / arg
+
+            with self.assertRaises(ZeroDivisionError):
+                compactlong_lhs(0)
+            with self.assertRaises(ZeroDivisionError):
+                compactlong_lhs(0.0)
+            with self.assertRaises(ZeroDivisionError):
+                float_lhs(0.0)
+            with self.assertRaises(ZeroDivisionError):
+                float_lhs(0)
+
+            self.assert_no_opcode(compactlong_lhs, "BINARY_OP_EXTEND")
+            self.assert_no_opcode(float_lhs, "BINARY_OP_EXTEND")
+
+        binary_op_zero_division()
+
+        def binary_op_nan():
+            def compactlong_lhs(arg):
+                42 + arg
+                42 - arg
+                42 * arg
+                42 / arg
+            def compactlong_rhs(arg):
+                arg + 42
+                arg - 42
+                arg * 42
+                arg / 42
+
+            compactlong_lhs(1.0)
+            compactlong_lhs(float('nan'))
+            compactlong_rhs(1.0)
+            compactlong_rhs(float('nan'))
+
+            self.assert_no_opcode(compactlong_lhs, "BINARY_OP_EXTEND")
+            self.assert_no_opcode(compactlong_rhs, "BINARY_OP_EXTEND")
+
+        binary_op_nan()
 
     @cpython_only
     @requires_specialization_ft

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2416,25 +2416,28 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs)
 
 /* float-long */
 
+// Guards should check that the float part is not NaN.
+// Guards for / (aka true_div) should check for 0 or 0.0 as the rhs.
 static int
 float_compactlong_guard(PyObject *lhs, PyObject *rhs)
 {
     return (
         PyFloat_CheckExact(lhs) &&
         PyLong_CheckExact(rhs) &&
-        _PyLong_IsCompact((PyLongObject *)rhs)
+        _PyLong_IsCompact((PyLongObject *)rhs) &&
+        !isnan(PyFloat_AsDouble(lhs))
     );
 }
 
 static int
 float_compactlong_guard_true_div(PyObject *lhs, PyObject *rhs)
 {
-    // guards should check if rhs has a non-zero value
     return (
         PyFloat_CheckExact(lhs) &&
         PyLong_CheckExact(rhs) &&
         _PyLong_IsCompact((PyLongObject *)rhs) &&
-        !PyLong_IsZero(rhs)
+        !PyLong_IsZero(rhs) &&
+        !isnan(PyFloat_AsDouble(lhs))
     );
 }
 
@@ -2460,19 +2463,20 @@ compactlong_float_guard(PyObject *lhs, PyObject *rhs)
     return (
         PyFloat_CheckExact(rhs) &&
         PyLong_CheckExact(lhs) &&
-        _PyLong_IsCompact((PyLongObject *)lhs)
+        _PyLong_IsCompact((PyLongObject *)lhs) &&
+        !isnan(PyFloat_AsDouble(rhs))
     );
 }
 
 static int
 compactlong_float_guard_true_div(PyObject *lhs, PyObject *rhs)
 {
-    // guards should check if rhs has a non-zero value
     return (
         PyFloat_CheckExact(rhs) &&
         PyLong_CheckExact(lhs) &&
         _PyLong_IsCompact((PyLongObject *)lhs) &&
-        PyFloat_AsDouble(rhs) != 0.0
+        PyFloat_AsDouble(rhs) != 0.0 &&
+        !isnan(PyFloat_AsDouble(rhs))
     );
 }