gh-104615: don't make unsafe swaps in apply_static_swaps (#104620)

2 files changed, 18 insertions, 12 deletions

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 5802f1d..3ee3323 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -616,16 +616,17 @@ _testinternalcapi.optimize_cfg -> object
 
   instructions: object
   consts: object
+  nlocals: int
 
 Apply compiler optimizations to an instruction list.
 [clinic start generated code]*/
 
 static PyObject *
 _testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions,
-                                    PyObject *consts)
-/*[clinic end generated code: output=5412aeafca683c8b input=7e8a3de86ebdd0f9]*/
+                                    PyObject *consts, int nlocals)
+/*[clinic end generated code: output=57c53c3a3dfd1df0 input=6a96d1926d58d7e5]*/
 {
-    return _PyCompile_OptimizeCfg(instructions, consts);
+    return _PyCompile_OptimizeCfg(instructions, consts, nlocals);
 }
 
 static int
diff --git a/Modules/clinic/_testinternalcapi.c.h b/Modules/clinic/_testinternalcapi.c.h
index 41dd504..f512412 100644
--- a/Modules/clinic/_testinternalcapi.c.h
+++ b/Modules/clinic/_testinternalcapi.c.h
@@ -83,7 +83,7 @@ exit:
 }
 
 PyDoc_STRVAR(_testinternalcapi_optimize_cfg__doc__,
-"optimize_cfg($module, /, instructions, consts)\n"
+"optimize_cfg($module, /, instructions, consts, nlocals)\n"
 "--\n"
 "\n"
 "Apply compiler optimizations to an instruction list.");
@@ -93,7 +93,7 @@ PyDoc_STRVAR(_testinternalcapi_optimize_cfg__doc__,
 
 static PyObject *
 _testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions,
-                                    PyObject *consts);
+                                    PyObject *consts, int nlocals);
 
 static PyObject *
 _testinternalcapi_optimize_cfg(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
@@ -101,14 +101,14 @@ _testinternalcapi_optimize_cfg(PyObject *module, PyObject *const *args, Py_ssize
     PyObject *return_value = NULL;
     #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
 
-    #define NUM_KEYWORDS 2
+    #define NUM_KEYWORDS 3
     static struct {
         PyGC_Head _this_is_not_used;
         PyObject_VAR_HEAD
         PyObject *ob_item[NUM_KEYWORDS];
     } _kwtuple = {
         .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
-        .ob_item = { &_Py_ID(instructions), &_Py_ID(consts), },
+        .ob_item = { &_Py_ID(instructions), &_Py_ID(consts), &_Py_ID(nlocals), },
     };
     #undef NUM_KEYWORDS
     #define KWTUPLE (&_kwtuple.ob_base.ob_base)
@@ -117,24 +117,29 @@ _testinternalcapi_optimize_cfg(PyObject *module, PyObject *const *args, Py_ssize
     #  define KWTUPLE NULL
     #endif  // !Py_BUILD_CORE
 
-    static const char * const _keywords[] = {"instructions", "consts", NULL};
+    static const char * const _keywords[] = {"instructions", "consts", "nlocals", NULL};
     static _PyArg_Parser _parser = {
         .keywords = _keywords,
         .fname = "optimize_cfg",
         .kwtuple = KWTUPLE,
     };
     #undef KWTUPLE
-    PyObject *argsbuf[2];
+    PyObject *argsbuf[3];
     PyObject *instructions;
     PyObject *consts;
+    int nlocals;
 
-    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 2, 2, 0, argsbuf);
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 3, 3, 0, argsbuf);
     if (!args) {
         goto exit;
     }
     instructions = args[0];
     consts = args[1];
-    return_value = _testinternalcapi_optimize_cfg_impl(module, instructions, consts);
+    nlocals = _PyLong_AsInt(args[2]);
+    if (nlocals == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = _testinternalcapi_optimize_cfg_impl(module, instructions, consts, nlocals);
 
 exit:
     return return_value;
@@ -201,4 +206,4 @@ _testinternalcapi_assemble_code_object(PyObject *module, PyObject *const *args,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=ab661d56a14b1a1c input=a9049054013a1b77]*/
+/*[clinic end generated code: output=2965f1578b986218 input=a9049054013a1b77]*/