gh-106149: Simplify stack depth calculation. Replace asserts by exceptions. (#107255)

diff --git a/Include/internal/pycore_flowgraph.h b/Include/internal/pycore_flowgraph.h
index 83b60ed849a950bcbf87af1cb8e5ccedc17c4a9f..cb55f8712e7ad60a9a64592d53b0ad540b769592 100644 (file)
--- a/Include/internal/pycore_flowgraph.h
+++ b/Include/internal/pycore_flowgraph.h
@@ -89,8 +89,8 @@ int _PyCfgBuilder_Init(_PyCfgBuilder *g);
 void _PyCfgBuilder_Fini(_PyCfgBuilder *g);
 
 int _PyCfg_OptimizeCodeUnit(_PyCfgBuilder *g, PyObject *consts, PyObject *const_cache,
-                            int code_flags, int nlocals, int nparams, int firstlineno);
-int _PyCfg_Stackdepth(_PyCfgBasicblock *entryblock, int code_flags);
+                            int nlocals, int nparams, int firstlineno);
+int _PyCfg_Stackdepth(_PyCfgBuilder *g);
 void _PyCfg_ConvertPseudoOps(_PyCfgBasicblock *entryblock);
 int _PyCfg_ResolveJumps(_PyCfgBuilder *g);
 

diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
index 82b0b50d0ea4374b331600a7a7df72f8a47f838b..fba41f0b119796cd21b21032b0c481cf2fd2fe7b 100644 (file)
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -991,6 +991,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('LOAD_NAME', 1, 11),
             ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
             ('LOAD_CONST', 2, 13),
+            ('RETURN_VALUE', 13),
             lbl,
             ('LOAD_CONST', 3, 14),
             ('RETURN_VALUE', 14),
@@ -998,7 +999,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
         expected_insts = [
             ('LOAD_NAME', 1, 11),
             ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
-            ('LOAD_CONST', 1, 13),
+            ('RETURN_CONST', 1, 13),
             lbl,
             ('RETURN_CONST', 2, 14),
         ]
@@ -1072,6 +1073,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('STORE_FAST', 1, 4),
             ('STORE_FAST', 1, 4),
             ('POP_TOP', 0, 4),
+            ('LOAD_CONST', 0, 5),
             ('RETURN_VALUE', 5)
         ]
         expected_insts = [
@@ -1080,7 +1082,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('NOP', 0, 3),
             ('STORE_FAST', 1, 4),
             ('POP_TOP', 0, 4),
-            ('RETURN_VALUE', 5)
+            ('RETURN_CONST', 0)
         ]
         self.cfg_optimization_test(insts, expected_insts, consts=list(range(3)), nlocals=1)
 
@@ -1092,6 +1094,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('STORE_FAST', 1, 4),
             ('STORE_FAST', 1, 4),
             ('STORE_FAST', 1, 4),
+            ('LOAD_CONST', 0, 5),
             ('RETURN_VALUE', 5)
         ]
         expected_insts = [
@@ -1100,7 +1103,7 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('NOP', 0, 3),
             ('POP_TOP', 0, 4),
             ('STORE_FAST', 1, 4),
-            ('RETURN_VALUE', 5)
+            ('RETURN_CONST', 0, 5)
         ]
         self.cfg_optimization_test(insts, expected_insts, consts=list(range(3)), nlocals=1)
 
@@ -1112,9 +1115,19 @@ class DirectCfgOptimizerTests(CfgOptimizationTestCase):
             ('STORE_FAST', 1, 4),
             ('STORE_FAST', 1, 5),
             ('STORE_FAST', 1, 6),
+            ('LOAD_CONST', 0, 5),
             ('RETURN_VALUE', 5)
         ]
-        self.cfg_optimization_test(insts, insts, consts=list(range(3)), nlocals=1)
+        expected_insts = [
+            ('LOAD_CONST', 0, 1),
+            ('LOAD_CONST', 1, 2),
+            ('LOAD_CONST', 2, 3),
+            ('STORE_FAST', 1, 4),
+            ('STORE_FAST', 1, 5),
+            ('STORE_FAST', 1, 6),
+            ('RETURN_CONST', 0, 5)
+        ]
+        self.cfg_optimization_test(insts, expected_insts, consts=list(range(3)), nlocals=1)
 
 
 if __name__ == "__main__":

diff --git a/Python/compile.c b/Python/compile.c
index 5915a4bc2f81b032d61debdc44d2bb2de7b0e45d..b4e06e7cb140851bc857092fdfc318aff461d865 100644 (file)
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7527,6 +7527,9 @@ build_cellfixedoffsets(_PyCompile_CodeUnitMetadata *umd)
     return fixed;
 }
 
+#define IS_GENERATOR(CF) \
+    ((CF) & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR))
+
 static int
 insert_prefix_instructions(_PyCompile_CodeUnitMetadata *umd, basicblock *entryblock,
                            int *fixed, int nfreevars, int code_flags)
@@ -7534,7 +7537,11 @@ insert_prefix_instructions(_PyCompile_CodeUnitMetadata *umd, basicblock *entrybl
     assert(umd->u_firstlineno > 0);
 
     /* Add the generator prefix instructions. */
-    if (code_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
+    if (IS_GENERATOR(code_flags)) {
+        /* Note that RETURN_GENERATOR + POP_TOP have a net stack effect
+         * of 0. This is because RETURN_GENERATOR pushes an element
+         * with _PyFrame_StackPush before switching stacks.
+         */
         cfg_instr make_gen = {
             .i_opcode = RETURN_GENERATOR,
             .i_oparg = 0,
@@ -7715,19 +7722,27 @@ optimize_and_assemble_code_unit(struct compiler_unit *u, PyObject *const_cache,
     int nparams = (int)PyList_GET_SIZE(u->u_ste->ste_varnames);
     int nlocals = (int)PyDict_GET_SIZE(u->u_metadata.u_varnames);
     assert(u->u_metadata.u_firstlineno);
-    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, code_flags, nlocals,
+    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, nlocals,
                                 nparams, u->u_metadata.u_firstlineno) < 0) {
         goto error;
     }
 
-    /** Assembly **/
-    int nlocalsplus = prepare_localsplus(&u->u_metadata, &g, code_flags);
-    if (nlocalsplus < 0) {
+    int stackdepth = _PyCfg_Stackdepth(&g);
+    if (stackdepth < 0) {
         goto error;
     }
 
-    int maxdepth = _PyCfg_Stackdepth(g.g_entryblock, code_flags);
-    if (maxdepth < 0) {
+    /* prepare_localsplus adds instructions for generators that push
+     * and pop an item on the stack. This assertion makes sure there
+     * is space on the stack for that.
+     * It should always be true, because at least one expression is
+     * required to turn a function into a generator.
+     */
+    assert(!(IS_GENERATOR(code_flags) && stackdepth == 0));
+
+    /** Assembly **/
+    int nlocalsplus = prepare_localsplus(&u->u_metadata, &g, code_flags);
+    if (nlocalsplus < 0) {
         goto error;
     }
 
@@ -7746,7 +7761,7 @@ optimize_and_assemble_code_unit(struct compiler_unit *u, PyObject *const_cache,
     }
 
     co = _PyAssemble_MakeCodeObject(&u->u_metadata, const_cache, consts,
-                                    maxdepth, &optimized_instrs, nlocalsplus,
+                                    stackdepth, &optimized_instrs, nlocalsplus,
                                     code_flags, filename);
 
 error:
@@ -8190,8 +8205,8 @@ _PyCompile_OptimizeCfg(PyObject *instructions, PyObject *consts, int nlocals)
     if (instructions_to_cfg(instructions, &g) < 0) {
         goto error;
     }
-    int code_flags = 0, nparams = 0, firstlineno = 1;
-    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, code_flags, nlocals,
+    int nparams = 0, firstlineno = 1;
+    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, nlocals,
                                 nparams, firstlineno) < 0) {
         goto error;
     }
@@ -8226,14 +8241,14 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
         goto error;
     }
 
-    int code_flags = 0;
-    int nlocalsplus = prepare_localsplus(umd, &g, code_flags);
-    if (nlocalsplus < 0) {
+    int stackdepth = _PyCfg_Stackdepth(&g);
+    if (stackdepth < 0) {
         goto error;
     }
 
-    int maxdepth = _PyCfg_Stackdepth(g.g_entryblock, code_flags);
-    if (maxdepth < 0) {
+    int code_flags = 0;
+    int nlocalsplus = prepare_localsplus(umd, &g, code_flags);
+    if (nlocalsplus < 0) {
         goto error;
     }
 
@@ -8256,7 +8271,7 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
         goto error;
     }
     co = _PyAssemble_MakeCodeObject(umd, const_cache,
-                                    consts, maxdepth, &optimized_instrs,
+                                    consts, stackdepth, &optimized_instrs,
                                     nlocalsplus, code_flags, filename);
     Py_DECREF(consts);
 

diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index f8e30ef7515a275f26b6013a8aaaafd3bee209ce..a72b85d45aa27dc4ace40ece9e22f9c604e8bcb8 100644 (file)
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -615,23 +615,28 @@ make_cfg_traversal_stack(basicblock *entryblock) {
     return stack;
 }
 
-Py_LOCAL_INLINE(void)
+Py_LOCAL_INLINE(int)
 stackdepth_push(basicblock ***sp, basicblock *b, int depth)
 {
-    assert(b->b_startdepth < 0 || b->b_startdepth == depth);
+    if (!(b->b_startdepth < 0 || b->b_startdepth == depth)) {
+        PyErr_Format(PyExc_ValueError, "Invalid CFG, inconsistent stackdepth");
+        return ERROR;
+    }
     if (b->b_startdepth < depth && b->b_startdepth < 100) {
         assert(b->b_startdepth < 0);
         b->b_startdepth = depth;
         *(*sp)++ = b;
     }
+    return SUCCESS;
 }
 
 /* Find the flow path that needs the largest stack.  We assume that
  * cycles in the flow graph have no net effect on the stack depth.
  */
 int
-_PyCfg_Stackdepth(basicblock *entryblock, int code_flags)
+_PyCfg_Stackdepth(cfg_builder *g)
 {
+    basicblock *entryblock = g->g_entryblock;
     for (basicblock *b = entryblock; b != NULL; b = b->b_next) {
         b->b_startdepth = INT_MIN;
     }
@@ -640,14 +645,13 @@ _PyCfg_Stackdepth(basicblock *entryblock, int code_flags)
         return ERROR;
     }
 
+
+    int stackdepth = -1;
     int maxdepth = 0;
     basicblock **sp = stack;
-    if (code_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
-        stackdepth_push(&sp, entryblock, 1);
-    } else {
-        stackdepth_push(&sp, entryblock, 0);
+    if (stackdepth_push(&sp, entryblock, 0) < 0) {
+        goto error;
     }
-
     while (sp != stack) {
         basicblock *b = *--sp;
         int depth = b->b_startdepth;
@@ -655,27 +659,40 @@ _PyCfg_Stackdepth(basicblock *entryblock, int code_flags)
         basicblock *next = b->b_next;
         for (int i = 0; i < b->b_iused; i++) {
             cfg_instr *instr = &b->b_instr[i];
-            int effect = PyCompile_OpcodeStackEffectWithJump(instr->i_opcode, instr->i_oparg, 0);
+            int effect = PyCompile_OpcodeStackEffectWithJump(
+                             instr->i_opcode, instr->i_oparg, 0);
             if (effect == PY_INVALID_STACK_EFFECT) {
                 PyErr_Format(PyExc_SystemError,
-                             "compiler PyCompile_OpcodeStackEffectWithJump(opcode=%d, arg=%i) failed",
+                             "Invalid stack effect for opcode=%d, arg=%i",
                              instr->i_opcode, instr->i_oparg);
-                return ERROR;
+                goto error;
             }
             int new_depth = depth + effect;
-            assert(new_depth >= 0); /* invalid code or bug in stackdepth() */
+            if (new_depth < 0) {
+               PyErr_Format(PyExc_ValueError,
+                            "Invalid CFG, stack underflow");
+               goto error;
+            }
             if (new_depth > maxdepth) {
                 maxdepth = new_depth;
             }
             if (HAS_TARGET(instr->i_opcode)) {
-                effect = PyCompile_OpcodeStackEffectWithJump(instr->i_opcode, instr->i_oparg, 1);
-                assert(effect != PY_INVALID_STACK_EFFECT);
+                effect = PyCompile_OpcodeStackEffectWithJump(
+                             instr->i_opcode, instr->i_oparg, 1);
+                if (effect == PY_INVALID_STACK_EFFECT) {
+                    PyErr_Format(PyExc_SystemError,
+                                 "Invalid stack effect for opcode=%d, arg=%i",
+                                 instr->i_opcode, instr->i_oparg);
+                    goto error;
+                }
                 int target_depth = depth + effect;
                 assert(target_depth >= 0); /* invalid code or bug in stackdepth() */
                 if (target_depth > maxdepth) {
                     maxdepth = target_depth;
                 }
-                stackdepth_push(&sp, instr->i_target, target_depth);
+                if (stackdepth_push(&sp, instr->i_target, target_depth) < 0) {
+                    goto error;
+                }
             }
             depth = new_depth;
             assert(!IS_ASSEMBLER_OPCODE(instr->i_opcode));
@@ -689,11 +706,15 @@ _PyCfg_Stackdepth(basicblock *entryblock, int code_flags)
         }
         if (next != NULL) {
             assert(BB_HAS_FALLTHROUGH(b));
-            stackdepth_push(&sp, next, depth);
+            if (stackdepth_push(&sp, next, depth) < 0) {
+                goto error;
+            }
         }
     }
+    stackdepth = maxdepth;
+error:
     PyMem_Free(stack);
-    return maxdepth;
+    return stackdepth;
 }
 
 static int
@@ -2009,7 +2030,7 @@ mark_cold(basicblock *entryblock) {
 
 
 static int
-push_cold_blocks_to_end(cfg_builder *g, int code_flags) {
+push_cold_blocks_to_end(cfg_builder *g) {
     basicblock *entryblock = g->g_entryblock;
     if (entryblock->b_next == NULL) {
         /* single basicblock, no need to reorder */
@@ -2251,7 +2272,7 @@ resolve_line_numbers(cfg_builder *g, int firstlineno)
 
 int
 _PyCfg_OptimizeCodeUnit(cfg_builder *g, PyObject *consts, PyObject *const_cache,
-                       int code_flags, int nlocals, int nparams, int firstlineno)
+                        int nlocals, int nparams, int firstlineno)
 {
     assert(cfg_builder_check(g));
     /** Preprocessing **/
@@ -2268,7 +2289,7 @@ _PyCfg_OptimizeCodeUnit(cfg_builder *g, PyObject *consts, PyObject *const_cache,
             g->g_entryblock, nlocals, nparams));
     insert_superinstructions(g);
 
-    RETURN_IF_ERROR(push_cold_blocks_to_end(g, code_flags));
+    RETURN_IF_ERROR(push_cold_blocks_to_end(g));
     RETURN_IF_ERROR(resolve_line_numbers(g, firstlineno));
     return SUCCESS;
 }