bpo-24340: Fix estimation of the code stack size. (#5076)

1 files changed, 90 insertions, 37 deletions

diff --git a/Python/compile.c b/Python/compile.c
index cd039168eb..7ba2ff34df 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -857,10 +857,22 @@ compiler_set_lineno(struct compiler *c, int off)
     b->b_instr[off].i_lineno = c->u->u_lineno;
 }
 
-int
-PyCompile_OpcodeStackEffect(int opcode, int oparg)
+/* Return the stack effect of opcode with argument oparg.
+
+   Some opcodes have different stack effect when jump to the target and
+   when not jump. The 'jump' parameter specifies the case:
+
+   * 0 -- when not jump
+   * 1 -- when jump
+   * -1 -- maximal
+ */
+/* XXX Make the stack effect of WITH_CLEANUP_START and
+   WITH_CLEANUP_FINISH deterministic. */
+static int
+stack_effect(int opcode, int oparg, int jump)
 {
     switch (opcode) {
+        /* Stack manipulation */
         case POP_TOP:
             return -1;
         case ROT_TWO:
@@ -871,6 +883,7 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case DUP_TOP_TWO:
             return 2;
 
+        /* Unary operators */
         case UNARY_POSITIVE:
         case UNARY_NEGATIVE:
         case UNARY_NOT:
@@ -883,6 +896,7 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case MAP_ADD:
             return -2;
 
+        /* Binary operators */
         case BINARY_POWER:
         case BINARY_MULTIPLY:
         case BINARY_MATRIX_MULTIPLY:
@@ -932,11 +946,16 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case BREAK_LOOP:
             return 0;
         case SETUP_WITH:
-            return 7;
+            /* 1 in the normal flow.
+             * Restore the stack position and push 6 values before jumping to
+             * the handler if an exception be raised. */
+            return jump ? 6 : 1;
         case WITH_CLEANUP_START:
-            return 1;
+            return 2; /* or 1, depending on TOS */
         case WITH_CLEANUP_FINISH:
-            return -1; /* XXX Sometimes more */
+            /* Pop a variable number of values pushed by WITH_CLEANUP_START
+             * + __exit__ or __aexit__. */
+            return -3;
         case RETURN_VALUE:
             return -1;
         case IMPORT_STAR:
@@ -950,9 +969,10 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case POP_BLOCK:
             return 0;
         case POP_EXCEPT:
-            return 0;  /* -3 except if bad bytecode */
+            return -3;
         case END_FINALLY:
-            return -1; /* or -2 or -3 if exception occurred */
+            /* Pop 6 values when an exception was raised. */
+            return -6;
 
         case STORE_NAME:
             return -1;
@@ -963,7 +983,8 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case UNPACK_EX:
             return (oparg&0xFF) + (oparg>>8);
         case FOR_ITER:
-            return 1; /* or -1, at end of iterator */
+            /* -1 at end of iterator, 1 if continue iterating. */
+            return jump > 0 ? -1 : 1;
 
         case STORE_ATTR:
             return -2;
@@ -1002,12 +1023,15 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
         case IMPORT_FROM:
             return 1;
 
+        /* Jumps */
         case JUMP_FORWARD:
-        case JUMP_IF_TRUE_OR_POP:  /* -1 if jump not taken */
-        case JUMP_IF_FALSE_OR_POP:  /*  "" */
         case JUMP_ABSOLUTE:
             return 0;
 
+        case JUMP_IF_TRUE_OR_POP:
+        case JUMP_IF_FALSE_OR_POP:
+            return jump ? 0 : -1;
+
         case POP_JUMP_IF_FALSE:
         case POP_JUMP_IF_TRUE:
             return -1;
@@ -1021,8 +1045,10 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
             return 0;
         case SETUP_EXCEPT:
         case SETUP_FINALLY:
-            return 6; /* can push 3 values for the new exception
-                + 3 others for the previous exception state */
+            /* 0 in the normal flow.
+             * Restore the stack position and push 6 values before jumping to
+             * the handler if an exception be raised. */
+            return jump ? 6 : 0;
 
         case LOAD_FAST:
             return 1;
@@ -1035,6 +1061,8 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
 
         case RAISE_VARARGS:
             return -oparg;
+
+        /* Functions and calls */
         case CALL_FUNCTION:
             return -oparg;
         case CALL_METHOD:
@@ -1052,6 +1080,7 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
             else
                 return -1;
 
+        /* Closures */
         case LOAD_CLOSURE:
             return 1;
         case LOAD_DEREF:
@@ -1061,10 +1090,16 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
             return -1;
         case DELETE_DEREF:
             return 0;
+
+        /* Iterators and generators */
         case GET_AWAITABLE:
             return 0;
         case SETUP_ASYNC_WITH:
-            return 6;
+            /* 0 in the normal flow.
+             * Restore the stack position to the position before the result
+             * of __aenter__ and push 6 values before jumping to the handler
+             * if an exception be raised. */
+            return jump ? -1 + 6 : 0;
         case BEFORE_ASYNC_WITH:
             return 1;
         case GET_AITER:
@@ -1085,6 +1120,12 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
     return PY_INVALID_STACK_EFFECT; /* not reachable */
 }
 
+int
+PyCompile_OpcodeStackEffect(int opcode, int oparg)
+{
+    return stack_effect(opcode, oparg, -1);
+}
+
 /* Add an opcode with no argument.
    Returns 0 on failure, 1 on success.
 */
@@ -2329,6 +2370,7 @@ compiler_async_for(struct compiler *c, stmt_ty s)
     ADDOP(c, POP_TOP);
     ADDOP(c, POP_TOP);
     ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
+    ADDOP(c, POP_TOP); /* for correct calculation of stack effect */
     ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
     ADDOP_JABS(c, JUMP_ABSOLUTE, after_loop_else);
 
@@ -2611,7 +2653,6 @@ compiler_try_except(struct compiler *c, stmt_ty s)
             /* second # body */
             VISIT_SEQ(c, stmt, handler->v.ExceptHandler.body);
             ADDOP(c, POP_BLOCK);
-            ADDOP(c, POP_EXCEPT);
             compiler_pop_fblock(c, FINALLY_TRY, cleanup_body);
 
             /* finally: */
@@ -2628,6 +2669,7 @@ compiler_try_except(struct compiler *c, stmt_ty s)
             compiler_nameop(c, handler->v.ExceptHandler.name, Del);
 
             ADDOP(c, END_FINALLY);
+            ADDOP(c, POP_EXCEPT);
             compiler_pop_fblock(c, FINALLY_END, cleanup_end);
         }
         else {
@@ -4899,44 +4941,55 @@ dfs(struct compiler *c, basicblock *b, struct assembler *a)
 static int
 stackdepth_walk(struct compiler *c, basicblock *b, int depth, int maxdepth)
 {
-    int i, target_depth, effect;
+    int i, new_depth, target_depth, effect;
     struct instr *instr;
-    if (b->b_seen || b->b_startdepth >= depth)
+    assert(!b->b_seen || b->b_startdepth == depth);
+    if (b->b_seen || b->b_startdepth >= depth) {
         return maxdepth;
+    }
+    /* Guard against infinite recursion */
     b->b_seen = 1;
     b->b_startdepth = depth;
     for (i = 0; i < b->b_iused; i++) {
         instr = &b->b_instr[i];
-        effect = PyCompile_OpcodeStackEffect(instr->i_opcode, instr->i_oparg);
+        effect = stack_effect(instr->i_opcode, instr->i_oparg, 0);
         if (effect == PY_INVALID_STACK_EFFECT) {
             fprintf(stderr, "opcode = %d\n", instr->i_opcode);
             Py_FatalError("PyCompile_OpcodeStackEffect()");
         }
-        depth += effect;
-
-        if (depth > maxdepth)
-            maxdepth = depth;
-        assert(depth >= 0); /* invalid code or bug in stackdepth() */
+        new_depth = depth + effect;
+        if (new_depth > maxdepth) {
+            maxdepth = new_depth;
+        }
+        assert(new_depth >= 0); /* invalid code or bug in stackdepth() */
         if (instr->i_jrel || instr->i_jabs) {
-            target_depth = depth;
-            if (instr->i_opcode == FOR_ITER) {
-                target_depth = depth-2;
+            /* Recursively inspect jump target */
+            effect = stack_effect(instr->i_opcode, instr->i_oparg, 1);
+            assert(effect != PY_INVALID_STACK_EFFECT);
+            target_depth = depth + effect;
+            if (target_depth > maxdepth) {
+                maxdepth = target_depth;
             }
-            else if (instr->i_opcode == SETUP_FINALLY ||
-                     instr->i_opcode == SETUP_EXCEPT) {
-                target_depth = depth+3;
-                if (target_depth > maxdepth)
-                    maxdepth = target_depth;
+            assert(target_depth >= 0); /* invalid code or bug in stackdepth() */
+            if (instr->i_opcode == CONTINUE_LOOP) {
+                /* Pops a variable number of values from the stack,
+                 * but the target should be already proceeding.
+                 */
+                assert(instr->i_target->b_seen);
+                assert(instr->i_target->b_startdepth <= depth);
+                goto out; /* remaining code is dead */
             }
-            else if (instr->i_opcode == JUMP_IF_TRUE_OR_POP ||
-                     instr->i_opcode == JUMP_IF_FALSE_OR_POP)
-                depth = depth - 1;
             maxdepth = stackdepth_walk(c, instr->i_target,
                                        target_depth, maxdepth);
-            if (instr->i_opcode == JUMP_ABSOLUTE ||
-                instr->i_opcode == JUMP_FORWARD) {
-                goto out; /* remaining code is dead */
-            }
+        }
+        depth = new_depth;
+        if (instr->i_opcode == JUMP_ABSOLUTE ||
+            instr->i_opcode == JUMP_FORWARD ||
+            instr->i_opcode == RETURN_VALUE ||
+            instr->i_opcode == RAISE_VARARGS ||
+            instr->i_opcode == BREAK_LOOP)
+        {
+            goto out; /* remaining code is dead */
         }
     }
     if (b->b_next)