gh-100239: Handle NaN and zero division in guards for `BINARY_OP_EXTEND` (#128963)

Co-authored-by: Tomas R. <tomas.roun8@gmail.com>
Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>

Lib/test/test_opcache.py

@@ -1362,6 +1362,53 @@ 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 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):
+                return (
+                    42 + arg,
+                    42 - arg,
+                    42 * arg,
+                    42 / arg,
+                )
+            def compactlong_rhs(arg):
+                return (
+                    arg + 42,
+                    arg - 42,
+                    arg * 2,
+                    arg / 42,
+                )
+            nan = float('nan')
+            self.assertEqual(compactlong_lhs(1.0), (43.0, 41.0, 42.0, 42.0))
+            for _ in range(100):
+                self.assertTrue(all(filter(lambda x: x is nan, compactlong_lhs(nan))))
+            self.assertEqual(compactlong_rhs(42.0), (84.0, 0.0, 84.0, 1.0))
+            for _ in range(100):
+                self.assertTrue(all(filter(lambda x: x is nan, compactlong_rhs(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

Python/specialize.c

@@ -2416,16 +2416,25 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs)
 
 /* float-long */
 
-static int
+static inline int
 float_compactlong_guard(PyObject *lhs, PyObject *rhs)
 {
     return (
         PyFloat_CheckExact(lhs) &&
+        !isnan(PyFloat_AsDouble(lhs)) &&
         PyLong_CheckExact(rhs) &&
         _PyLong_IsCompact((PyLongObject *)rhs)
     );
 }
 
+static inline int
+nonzero_float_compactlong_guard(PyObject *lhs, PyObject *rhs)
+{
+    return (
+        float_compactlong_guard(lhs, rhs) && !PyLong_IsZero(rhs)
+    );
+}
+
 #define FLOAT_LONG_ACTION(NAME, OP) \
     static PyObject * \
     (NAME)(PyObject *lhs, PyObject *rhs) \
@@ -2442,13 +2451,22 @@ FLOAT_LONG_ACTION(float_compactlong_true_div, /)
 
 /*  long-float */
 
-static int
+static inline int
 compactlong_float_guard(PyObject *lhs, PyObject *rhs)
 {
     return (
-        PyFloat_CheckExact(rhs) &&
         PyLong_CheckExact(lhs) &&
-        _PyLong_IsCompact((PyLongObject *)lhs)
+        _PyLong_IsCompact((PyLongObject *)lhs) &&
+        PyFloat_CheckExact(rhs) &&
+        !isnan(PyFloat_AsDouble(rhs))
+    );
+}
+
+static inline int
+nonzero_compactlong_float_guard(PyObject *lhs, PyObject *rhs)
+{
+    return (
+        compactlong_float_guard(lhs, rhs) && PyFloat_AsDouble(rhs) != 0.0
     );
 }
 
@@ -2469,14 +2487,14 @@ LONG_FLOAT_ACTION(compactlong_float_true_div, /)
 static _PyBinaryOpSpecializationDescr float_compactlong_specs[NB_OPARG_LAST+1] = {
     [NB_ADD] = {float_compactlong_guard, float_compactlong_add},
     [NB_SUBTRACT] = {float_compactlong_guard, float_compactlong_subtract},
-    [NB_TRUE_DIVIDE] = {float_compactlong_guard, float_compactlong_true_div},
+    [NB_TRUE_DIVIDE] = {nonzero_float_compactlong_guard, float_compactlong_true_div},
     [NB_MULTIPLY] = {float_compactlong_guard, float_compactlong_multiply},
 };
 
 static _PyBinaryOpSpecializationDescr compactlong_float_specs[NB_OPARG_LAST+1] = {
     [NB_ADD] = {compactlong_float_guard, compactlong_float_add},
     [NB_SUBTRACT] = {compactlong_float_guard, compactlong_float_subtract},
-    [NB_TRUE_DIVIDE] = {compactlong_float_guard, compactlong_float_true_div},
+    [NB_TRUE_DIVIDE] = {nonzero_compactlong_float_guard, compactlong_float_true_div},
     [NB_MULTIPLY] = {compactlong_float_guard, compactlong_float_multiply},
 };