GH-90043: Handle NaNs in COMPARE_OP_FLOAT_JUMP (GH-100278)

3 files changed, 21 insertions, 24 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index c56f1d3..8e95b73 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -2024,13 +2024,11 @@ dummy_func(
             // Combined: COMPARE_OP (float ? float) + POP_JUMP_IF_(true/false)
             DEOPT_IF(!PyFloat_CheckExact(left), COMPARE_OP);
             DEOPT_IF(!PyFloat_CheckExact(right), COMPARE_OP);
+            STAT_INC(COMPARE_OP, hit);
             double dleft = PyFloat_AS_DOUBLE(left);
             double dright = PyFloat_AS_DOUBLE(right);
-            // 1 if <, 2 if ==, 4 if >; this matches when _to_jump_mask
-            int sign_ish = 2*(dleft > dright) + 2 - (dleft < dright);
-            DEOPT_IF(isnan(dleft), COMPARE_OP);
-            DEOPT_IF(isnan(dright), COMPARE_OP);
-            STAT_INC(COMPARE_OP, hit);
+            // 1 if NaN, 2 if <, 4 if >, 8 if ==; this matches when_to_jump_mask
+            int sign_ish = 1 << (2 * (dleft >= dright) + (dleft <= dright));
             _Py_DECREF_SPECIALIZED(left, _PyFloat_ExactDealloc);
             _Py_DECREF_SPECIALIZED(right, _PyFloat_ExactDealloc);
             jump = sign_ish & when_to_jump_mask;
@@ -2057,8 +2055,8 @@ dummy_func(
             assert(Py_ABS(Py_SIZE(left)) <= 1 && Py_ABS(Py_SIZE(right)) <= 1);
             Py_ssize_t ileft = Py_SIZE(left) * ((PyLongObject *)left)->ob_digit[0];
             Py_ssize_t iright = Py_SIZE(right) * ((PyLongObject *)right)->ob_digit[0];
-            // 1 if <, 2 if ==, 4 if >; this matches when _to_jump_mask
-            int sign_ish = 2*(ileft > iright) + 2 - (ileft < iright);
+            // 2 if <, 4 if >, 8 if ==; this matches when_to_jump_mask
+            int sign_ish = 1 << (2 * (ileft >= iright) + (ileft <= iright));
             _Py_DECREF_SPECIALIZED(left, (destructor)PyObject_Free);
             _Py_DECREF_SPECIALIZED(right, (destructor)PyObject_Free);
             jump = sign_ish & when_to_jump_mask;
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 63635fb..6d84a64 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -2214,13 +2214,11 @@
                 // Combined: COMPARE_OP (float ? float) + POP_JUMP_IF_(true/false)
                 DEOPT_IF(!PyFloat_CheckExact(left), COMPARE_OP);
                 DEOPT_IF(!PyFloat_CheckExact(right), COMPARE_OP);
+                STAT_INC(COMPARE_OP, hit);
                 double dleft = PyFloat_AS_DOUBLE(left);
                 double dright = PyFloat_AS_DOUBLE(right);
-                // 1 if <, 2 if ==, 4 if >; this matches when _to_jump_mask
-                int sign_ish = 2*(dleft > dright) + 2 - (dleft < dright);
-                DEOPT_IF(isnan(dleft), COMPARE_OP);
-                DEOPT_IF(isnan(dright), COMPARE_OP);
-                STAT_INC(COMPARE_OP, hit);
+                // 1 if NaN, 2 if <, 4 if >, 8 if ==; this matches when_to_jump_mask
+                int sign_ish = 1 << (2 * (dleft >= dright) + (dleft <= dright));
                 _Py_DECREF_SPECIALIZED(left, _PyFloat_ExactDealloc);
                 _Py_DECREF_SPECIALIZED(right, _PyFloat_ExactDealloc);
                 jump = sign_ish & when_to_jump_mask;
@@ -2258,8 +2256,8 @@
                 assert(Py_ABS(Py_SIZE(left)) <= 1 && Py_ABS(Py_SIZE(right)) <= 1);
                 Py_ssize_t ileft = Py_SIZE(left) * ((PyLongObject *)left)->ob_digit[0];
                 Py_ssize_t iright = Py_SIZE(right) * ((PyLongObject *)right)->ob_digit[0];
-                // 1 if <, 2 if ==, 4 if >; this matches when _to_jump_mask
-                int sign_ish = 2*(ileft > iright) + 2 - (ileft < iright);
+                // 2 if <, 4 if >, 8 if ==; this matches when_to_jump_mask
+                int sign_ish = 1 << (2 * (ileft >= iright) + (ileft <= iright));
                 _Py_DECREF_SPECIALIZED(left, (destructor)PyObject_Free);
                 _Py_DECREF_SPECIALIZED(right, (destructor)PyObject_Free);
                 jump = sign_ish & when_to_jump_mask;
diff --git a/Python/specialize.c b/Python/specialize.c
index 020127e..a1666cc 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -1951,15 +1951,16 @@ compare_op_fail_kind(PyObject *lhs, PyObject *rhs)
 
 
 static int compare_masks[] = {
-    // 1-bit: jump if less than
-    // 2-bit: jump if equal
+    // 1-bit: jump if unordered
+    // 2-bit: jump if less
     // 4-bit: jump if greater
-    [Py_LT] = 1 | 0 | 0,
-    [Py_LE] = 1 | 2 | 0,
-    [Py_EQ] = 0 | 2 | 0,
-    [Py_NE] = 1 | 0 | 4,
-    [Py_GT] = 0 | 0 | 4,
-    [Py_GE] = 0 | 2 | 4,
+    // 8-bit: jump if equal
+    [Py_LT] = 0 | 2 | 0 | 0,
+    [Py_LE] = 0 | 2 | 0 | 8,
+    [Py_EQ] = 0 | 0 | 0 | 8,
+    [Py_NE] = 1 | 2 | 4 | 0,
+    [Py_GT] = 0 | 0 | 4 | 0,
+    [Py_GE] = 0 | 0 | 4 | 8,
 };
 
 void
@@ -1980,7 +1981,7 @@ _Py_Specialize_CompareOp(PyObject *lhs, PyObject *rhs, _Py_CODEUNIT *instr,
     assert(oparg <= Py_GE);
     int when_to_jump_mask = compare_masks[oparg];
     if (next_opcode == POP_JUMP_IF_FALSE) {
-        when_to_jump_mask = (1 | 2 | 4) & ~when_to_jump_mask;
+        when_to_jump_mask = (1 | 2 | 4 | 8) & ~when_to_jump_mask;
     }
     if (Py_TYPE(lhs) != Py_TYPE(rhs)) {
         SPECIALIZATION_FAIL(COMPARE_OP, compare_op_fail_kind(lhs, rhs));
@@ -2009,7 +2010,7 @@ _Py_Specialize_CompareOp(PyObject *lhs, PyObject *rhs, _Py_CODEUNIT *instr,
         }
         else {
             _py_set_opcode(instr, COMPARE_OP_STR_JUMP);
-            cache->mask = (when_to_jump_mask & 2) == 0;
+            cache->mask = (when_to_jump_mask & 8) == 0;
             goto success;
         }
     }