gh-98831: Modernize CALL_FUNCTION_EX (#101627)

New generator feature: Move CHECK_EVAL_BREAKER() call to just before DISPATCH().

5 files changed, 41 insertions, 37 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index d0f0513a36f..9633f34212a 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -2951,26 +2951,21 @@ dummy_func(
             CHECK_EVAL_BREAKER();
         }
 
-        // error: CALL_FUNCTION_EX has irregular stack effect
-        inst(CALL_FUNCTION_EX) {
-            PyObject *func, *callargs, *kwargs = NULL, *result;
-            if (oparg & 0x01) {
-                kwargs = POP();
+        inst(CALL_FUNCTION_EX, (unused, func, callargs, kwargs if (oparg & 1) -- result)) {
+            if (oparg & 1) {
                 // DICT_MERGE is called before this opcode if there are kwargs.
                 // It converts all dict subtypes in kwargs into regular dicts.
                 assert(PyDict_CheckExact(kwargs));
             }
-            callargs = POP();
-            func = TOP();
             if (!PyTuple_CheckExact(callargs)) {
                 if (check_args_iterable(tstate, func, callargs) < 0) {
-                    Py_DECREF(callargs);
                     goto error;
                 }
-                Py_SETREF(callargs, PySequence_Tuple(callargs));
-                if (callargs == NULL) {
+                PyObject *tuple = PySequence_Tuple(callargs);
+                if (tuple == NULL) {
                     goto error;
                 }
+                Py_SETREF(callargs, tuple);
             }
             assert(PyTuple_CheckExact(callargs));
 
@@ -2979,12 +2974,8 @@ dummy_func(
             Py_DECREF(callargs);
             Py_XDECREF(kwargs);
 
-            STACK_SHRINK(1);
-            assert(TOP() == NULL);
-            SET_TOP(result);
-            if (result == NULL) {
-                goto error;
-            }
+            assert(PEEK(3 + (oparg & 1)) == NULL);
+            ERROR_IF(result == NULL, error);
             CHECK_EVAL_BREAKER();
         }
 
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 3ef808691e0..f38286441be 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -3587,24 +3587,24 @@
 
         TARGET(CALL_FUNCTION_EX) {
             PREDICTED(CALL_FUNCTION_EX);
-            PyObject *func, *callargs, *kwargs = NULL, *result;
-            if (oparg & 0x01) {
-                kwargs = POP();
+            PyObject *kwargs = (oparg & 1) ? PEEK(((oparg & 1) ? 1 : 0)) : NULL;
+            PyObject *callargs = PEEK(1 + ((oparg & 1) ? 1 : 0));
+            PyObject *func = PEEK(2 + ((oparg & 1) ? 1 : 0));
+            PyObject *result;
+            if (oparg & 1) {
                 // DICT_MERGE is called before this opcode if there are kwargs.
                 // It converts all dict subtypes in kwargs into regular dicts.
                 assert(PyDict_CheckExact(kwargs));
             }
-            callargs = POP();
-            func = TOP();
             if (!PyTuple_CheckExact(callargs)) {
                 if (check_args_iterable(tstate, func, callargs) < 0) {
-                    Py_DECREF(callargs);
                     goto error;
                 }
-                Py_SETREF(callargs, PySequence_Tuple(callargs));
-                if (callargs == NULL) {
+                PyObject *tuple = PySequence_Tuple(callargs);
+                if (tuple == NULL) {
                     goto error;
                 }
+                Py_SETREF(callargs, tuple);
             }
             assert(PyTuple_CheckExact(callargs));
 
@@ -3613,12 +3613,11 @@
             Py_DECREF(callargs);
             Py_XDECREF(kwargs);
 
-            STACK_SHRINK(1);
-            assert(TOP() == NULL);
-            SET_TOP(result);
-            if (result == NULL) {
-                goto error;
-            }
+            assert(PEEK(3 + (oparg & 1)) == NULL);
+            if (result == NULL) { STACK_SHRINK(((oparg & 1) ? 1 : 0)); goto pop_3_error; }
+            STACK_SHRINK(((oparg & 1) ? 1 : 0));
+            STACK_SHRINK(2);
+            POKE(1, result);
             CHECK_EVAL_BREAKER();
             DISPATCH();
         }
diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
index 52bab1c680e..054ef6c2998 100644
--- a/Python/opcode_metadata.h
+++ b/Python/opcode_metadata.h
@@ -325,7 +325,7 @@ _PyOpcode_num_popped(int opcode, int oparg, bool jump) {
         case CALL_NO_KW_METHOD_DESCRIPTOR_FAST:
             return -1;
         case CALL_FUNCTION_EX:
-            return -1;
+            return ((oparg & 1) ? 1 : 0) + 3;
         case MAKE_FUNCTION:
             return ((oparg & 0x01) ? 1 : 0) + ((oparg & 0x02) ? 1 : 0) + ((oparg & 0x04) ? 1 : 0) + ((oparg & 0x08) ? 1 : 0) + 1;
         case RETURN_GENERATOR:
@@ -673,7 +673,7 @@ _PyOpcode_num_pushed(int opcode, int oparg, bool jump) {
         case CALL_NO_KW_METHOD_DESCRIPTOR_FAST:
             return -1;
         case CALL_FUNCTION_EX:
-            return -1;
+            return 1;
         case MAKE_FUNCTION:
             return 1;
         case RETURN_GENERATOR:
diff --git a/Tools/cases_generator/generate_cases.py b/Tools/cases_generator/generate_cases.py
index 4f94b48d114..9b5aa914cde 100644
--- a/Tools/cases_generator/generate_cases.py
+++ b/Tools/cases_generator/generate_cases.py
@@ -227,7 +227,8 @@ class Instruction:
         self.kind = inst.kind
         self.name = inst.name
         self.block = inst.block
-        self.block_text, self.predictions = extract_block_text(self.block)
+        self.block_text, self.check_eval_breaker, self.predictions = \
+            extract_block_text(self.block)
         self.always_exits = always_exits(self.block_text)
         self.cache_effects = [
             effect for effect in inst.inputs if isinstance(effect, parser.CacheEffect)
@@ -1027,6 +1028,8 @@ class Analyzer:
             if not instr.always_exits:
                 for prediction in instr.predictions:
                     self.out.emit(f"PREDICT({prediction});")
+                if instr.check_eval_breaker:
+                    self.out.emit("CHECK_EVAL_BREAKER();")
                 self.out.emit(f"DISPATCH();")
 
     def write_super(self, sup: SuperInstruction) -> None:
@@ -1102,7 +1105,7 @@ class Analyzer:
             self.out.emit(f"DISPATCH();")
 
 
-def extract_block_text(block: parser.Block) -> tuple[list[str], list[str]]:
+def extract_block_text(block: parser.Block) -> tuple[list[str], bool, list[str]]:
     # Get lines of text with proper dedent
     blocklines = block.text.splitlines(True)
 
@@ -1122,6 +1125,12 @@ def extract_block_text(block: parser.Block) -> tuple[list[str], list[str]]:
     while blocklines and not blocklines[-1].strip():
         blocklines.pop()
 
+    # Separate CHECK_EVAL_BREAKER() macro from end
+    check_eval_breaker = \
+        blocklines != [] and blocklines[-1].strip() == "CHECK_EVAL_BREAKER();"
+    if check_eval_breaker:
+        del blocklines[-1]
+
     # Separate PREDICT(...) macros from end
     predictions: list[str] = []
     while blocklines and (
@@ -1130,7 +1139,7 @@ def extract_block_text(block: parser.Block) -> tuple[list[str], list[str]]:
         predictions.insert(0, m.group(1))
         blocklines.pop()
 
-    return blocklines, predictions
+    return blocklines, check_eval_breaker, predictions
 
 
 def always_exits(lines: list[str]) -> bool:
diff --git a/Tools/cases_generator/test_generator.py b/Tools/cases_generator/test_generator.py
index 33bba7ee340..036094ac8ef 100644
--- a/Tools/cases_generator/test_generator.py
+++ b/Tools/cases_generator/test_generator.py
@@ -177,15 +177,16 @@ def test_overlap():
     """
     run_cases_test(input, output)
 
-def test_predictions():
+def test_predictions_and_eval_breaker():
     input = """
         inst(OP1, (--)) {
         }
         inst(OP2, (--)) {
         }
-        inst(OP3, (--)) {
+        inst(OP3, (arg -- res)) {
             DEOPT_IF(xxx, OP1);
             PREDICT(OP2);
+            CHECK_EVAL_BREAKER();
         }
     """
     output = """
@@ -200,8 +201,12 @@ def test_predictions():
         }
 
         TARGET(OP3) {
+            PyObject *arg = PEEK(1);
+            PyObject *res;
             DEOPT_IF(xxx, OP1);
+            POKE(1, res);
             PREDICT(OP2);
+            CHECK_EVAL_BREAKER();
             DISPATCH();
         }
     """