diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index e0cc010f15513b..abb071451d8b59 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -1,3 +1,4 @@ +import ast import dis import gc from itertools import combinations, product @@ -1131,6 +1132,53 @@ def f(self): class DirectCfgOptimizerTests(CfgOptimizationTestCase): + def test_optimize_cfg_const_index_out_of_range(self): + insts = [ + ('LOAD_CONST', 2, 0), + ('RETURN_VALUE', None, 0), + ] + seq = self.seq_from_insts(insts) + with self.assertRaisesRegex(ValueError, "out of range"): + _testinternalcapi.optimize_cfg(seq, [0, 1], 0) + + def test_optimize_cfg_consts_must_be_list(self): + insts = [ + ('LOAD_CONST', 0, 0), + ('RETURN_VALUE', None, 0), + ] + seq = self.seq_from_insts(insts) + with self.assertRaisesRegex(TypeError, "consts must be a list"): + _testinternalcapi.optimize_cfg(seq, (0,), 0) + + def test_compiler_codegen_metadata_consts_roundtrips_optimize_cfg(self): + tree = ast.parse("x = (1, 2)", mode="exec", optimize=1) + insts, meta = _testinternalcapi.compiler_codegen(tree, "", 0) + consts = meta["consts"] + self.assertIsInstance(consts, list) + _testinternalcapi.optimize_cfg(insts, consts, 0) + + def test_compiler_codegen_consts_include_none_required_for_implicit_return(self): + # Module "pass" only needs the const table entry for None once + # _PyCodegen_AddReturnAtEnd runs. If metadata["consts"] were taken + # before that, the list would not match LOAD_CONST opargs (here: 0 + # for None), and optimize_cfg would read out of range. + tree = ast.parse("pass", mode="exec", optimize=1) + insts, meta = _testinternalcapi.compiler_codegen(tree, "", 0) + consts = meta["consts"] + self.assertEqual(consts, [None]) + + load_const = opcode.opmap["LOAD_CONST"] + self.assertEqual( + [t[1] for t in insts.get_instructions() if t[0] == load_const], + [0], + ) + + # As if consts were snapshotted before AddReturnAtEnd: still LOAD_CONST 0, no row. + with self.assertRaisesRegex(ValueError, "out of range"): + _testinternalcapi.optimize_cfg(insts, [], 0) + + _testinternalcapi.optimize_cfg(insts, list(consts), 0) + def cfg_optimization_test(self, insts, expected_insts, consts=None, expected_consts=None, nlocals=0): diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c index deac8570fe3241..619f9f50574429 100644 --- a/Modules/_testinternalcapi.c +++ b/Modules/_testinternalcapi.c @@ -1081,13 +1081,17 @@ _testinternalcapi.compiler_codegen -> object compile_mode: int = 0 Apply compiler code generation to an AST. + +Return (instruction_sequence, metadata). metadata maps "argcount", +"posonlyargcount", "kwonlyargcount" to ints and "consts" to the list of +constants in LOAD_CONST index order (for use with optimize_cfg). [clinic start generated code]*/ static PyObject * _testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast, PyObject *filename, int optimize, int compile_mode) -/*[clinic end generated code: output=40a68f6e13951cc8 input=a0e00784f1517cd7]*/ +/*[clinic end generated code: output=40a68f6e13951cc8 input=e0c65e5c80efe30e]*/ { PyCompilerFlags *flags = NULL; return _PyCompile_CodeGen(ast, filename, flags, optimize, compile_mode); @@ -1103,12 +1107,15 @@ _testinternalcapi.optimize_cfg -> object nlocals: int Apply compiler optimizations to an instruction list. + +consts must be a list aligned with LOAD_CONST opargs (the "consts" entry +from the metadata dict returned by compiler_codegen for the same unit). [clinic start generated code]*/ static PyObject * _testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions, PyObject *consts, int nlocals) -/*[clinic end generated code: output=57c53c3a3dfd1df0 input=6a96d1926d58d7e5]*/ +/*[clinic end generated code: output=57c53c3a3dfd1df0 input=905c3d935e063b27]*/ { return _PyCompile_OptimizeCfg(instructions, consts, nlocals); } diff --git a/Modules/clinic/_testinternalcapi.c.h b/Modules/clinic/_testinternalcapi.c.h index 21f4ee3201e5bf..85edc6fbb5802f 100644 --- a/Modules/clinic/_testinternalcapi.c.h +++ b/Modules/clinic/_testinternalcapi.c.h @@ -92,7 +92,11 @@ PyDoc_STRVAR(_testinternalcapi_compiler_codegen__doc__, "compiler_codegen($module, /, ast, filename, optimize, compile_mode=0)\n" "--\n" "\n" -"Apply compiler code generation to an AST."); +"Apply compiler code generation to an AST.\n" +"\n" +"Return (instruction_sequence, metadata). metadata maps \"argcount\",\n" +"\"posonlyargcount\", \"kwonlyargcount\" to ints and \"consts\" to the list of\n" +"constants in LOAD_CONST index order (for use with optimize_cfg)."); #define _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF \ {"compiler_codegen", _PyCFunction_CAST(_testinternalcapi_compiler_codegen), METH_FASTCALL|METH_KEYWORDS, _testinternalcapi_compiler_codegen__doc__}, @@ -169,7 +173,10 @@ PyDoc_STRVAR(_testinternalcapi_optimize_cfg__doc__, "optimize_cfg($module, /, instructions, consts, nlocals)\n" "--\n" "\n" -"Apply compiler optimizations to an instruction list."); +"Apply compiler optimizations to an instruction list.\n" +"\n" +"consts must be a list aligned with LOAD_CONST opargs (the \"consts\" entry\n" +"from the metadata dict returned by compiler_codegen for the same unit)."); #define _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF \ {"optimize_cfg", _PyCFunction_CAST(_testinternalcapi_optimize_cfg), METH_FASTCALL|METH_KEYWORDS, _testinternalcapi_optimize_cfg__doc__}, @@ -392,4 +399,4 @@ get_next_dict_keys_version(PyObject *module, PyObject *Py_UNUSED(ignored)) { return get_next_dict_keys_version_impl(module); } -/*[clinic end generated code: output=fbd8b7e0cae8bac7 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=ecb5d7ac85b153fa input=a9049054013a1b77]*/ diff --git a/Python/compile.c b/Python/compile.c index 5f82641a3948c6..eb9fc827bea40a 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -1658,6 +1658,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags, { PyObject *res = NULL; PyObject *metadata = NULL; + PyObject *consts_list = NULL; if (!PyAST_Check(ast)) { PyErr_SetString(PyExc_TypeError, "expected an AST"); @@ -1712,12 +1713,23 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags, } if (_PyInstructionSequence_ApplyLabelMap(_PyCompile_InstrSequence(c)) < 0) { - return NULL; + goto finally; + } + + /* After AddReturnAtEnd: co_consts indices match the final instruction stream. */ + consts_list = consts_dict_keys_inorder(umd->u_consts); + if (consts_list == NULL) { + goto finally; + } + if (PyDict_SetItemString(metadata, "consts", consts_list) < 0) { + goto finally; } + /* Allocate a copy of the instruction sequence on the heap */ res = _PyTuple_FromPair((PyObject *)_PyCompile_InstrSequence(c), metadata); finally: + Py_XDECREF(consts_list); Py_XDECREF(metadata); _PyCompile_ExitScope(c); compiler_free(c); diff --git a/Python/flowgraph.c b/Python/flowgraph.c index c234fa3d8c3876..202e3bacf2e1bf 100644 --- a/Python/flowgraph.c +++ b/Python/flowgraph.c @@ -1309,6 +1309,14 @@ get_const_value(int opcode, int oparg, PyObject *co_consts) PyObject *constant = NULL; assert(loads_const(opcode)); if (opcode == LOAD_CONST) { + assert(PyList_Check(co_consts)); + Py_ssize_t n = PyList_GET_SIZE(co_consts); + if (oparg < 0 || oparg >= n) { + PyErr_Format(PyExc_ValueError, + "LOAD_CONST index %d is out of range for consts (len=%zd)", + oparg, n); + return NULL; + } constant = PyList_GET_ITEM(co_consts, oparg); } if (opcode == LOAD_SMALL_INT) { @@ -2167,6 +2175,9 @@ basicblock_optimize_load_const(PyObject *const_cache, basicblock *bb, PyObject * cfg_instr *inst = &bb->b_instr[i]; if (inst->i_opcode == LOAD_CONST) { PyObject *constant = get_const_value(inst->i_opcode, inst->i_oparg, consts); + if (constant == NULL) { + return ERROR; + } int res = maybe_instr_make_load_smallint(inst, constant, consts, const_cache); Py_DECREF(constant); if (res < 0) { @@ -4064,6 +4075,10 @@ _PyCompile_OptimizeCfg(PyObject *seq, PyObject *consts, int nlocals) PyErr_SetString(PyExc_ValueError, "expected an instruction sequence"); return NULL; } + if (!PyList_Check(consts)) { + PyErr_SetString(PyExc_TypeError, "consts must be a list"); + return NULL; + } PyObject *const_cache = PyDict_New(); if (const_cache == NULL) { return NULL;