Factor out all block pops and pushes into simple opcodes.

python · markshannon · May 7, 2021 · Apr 19, 2021 · Apr 21, 2021 · Apr 21, 2021
commit 4c21ddd190521036bcf71ed579fc9f49f07b74d9
diff --git a/Include/opcode.h b/Include/opcode.h
diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -93,6 +93,7 @@ def jabs_op(name, op):
 def_op('GET_AITER', 50)
 def_op('GET_ANEXT', 51)
 def_op('BEFORE_ASYNC_WITH', 52)
+def_op('BEFORE_WITH', 53)
 
 def_op('END_ASYNC_FOR', 54)
 def_op('INPLACE_ADD', 55)
@@ -193,7 +194,7 @@ def jabs_op(name, op):
 
 def_op('CALL_FUNCTION_KW', 141)  # #args + #kwargs
 def_op('CALL_FUNCTION_EX', 142)  # Flags
-jrel_op('SETUP_WITH', 143)
+
 def_op('EXTENDED_ARG', 144)
 EXTENDED_ARG = 144
 def_op('LIST_APPEND', 145)

@@ -171,7 +171,6 @@ markblocks(PyCodeObject *code_obj, int len)
                     block_stack = push_block(block_stack, Try);
                     blocks[i+1] = block_stack;
                     break;
-                case SETUP_WITH:
                 case SETUP_ASYNC_WITH:
                     j = get_arg(code, i) + i + 1;
                     assert(j < len);

@@ -1728,7 +1728,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             opcode = _Py_OPCODE(*next_instr);
             if (opcode != SETUP_FINALLY &&
                 opcode != SETUP_CLEANUP &&
-                opcode != SETUP_WITH &&
                 opcode != BEFORE_ASYNC_WITH &&
                 opcode != YIELD_FROM) {
                 /* Few cases where we skip running signal handlers and other
@@ -2637,9 +2636,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
         case TARGET(POP_EXCEPT): {
             PyObject *type, *value, *traceback;
             _PyErr_StackItem *exc_info;
-            PyTryBlock *b = PyFrame_BlockPop(f);
-            assert(STACK_LEVEL() == (b)->b_level ||
-                   STACK_LEVEL() == (b)->b_level + 1);
             exc_info = tstate->exc_info;
             type = exc_info->exc_type;
             value = exc_info->exc_value;
@@ -4086,6 +4082,15 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
+        case TARGET(SETUP_ASYNC_WITH): {
+            PREDICTED(SETUP_ASYNC_WITH);
+            /* Setup the finally block before pushing the result
+               of __aenter__ on the stack. */
+            PyFrame_BlockSetup(f, SETUP_CLEANUP, INSTR_OFFSET() + oparg,
+                               STACK_LEVEL()-1);
+            DISPATCH();
+        }
+
         case TARGET(BEFORE_ASYNC_WITH): {
             _Py_IDENTIFIER(__aenter__);
             _Py_IDENTIFIER(__aexit__);
@@ -4111,17 +4116,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
-        case TARGET(SETUP_ASYNC_WITH): {
-            PyObject *res = POP();
-            /* Setup the finally block before pushing the result
-               of __aenter__ on the stack. */
-            PyFrame_BlockSetup(f, SETUP_CLEANUP, INSTR_OFFSET() + oparg,
-                               STACK_LEVEL());
-            PUSH(res);
-            DISPATCH();
-        }
-
-        case TARGET(SETUP_WITH): {
+        case TARGET(BEFORE_WITH): {
             _Py_IDENTIFIER(__enter__);
             _Py_IDENTIFIER(__exit__);
             PyObject *mgr = TOP();
@@ -4139,14 +4134,11 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             Py_DECREF(mgr);
             res = _PyObject_CallNoArg(enter);
             Py_DECREF(enter);
-            if (res == NULL)
+            if (res == NULL) {
                 goto error;
-            /* Setup the finally block before pushing the result
-               of __enter__ on the stack. */
-            PyFrame_BlockSetup(f, SETUP_CLEANUP, INSTR_OFFSET() + oparg,
-                               STACK_LEVEL());
-
+            }
             PUSH(res);
+            PREDICT(SETUP_ASYNC_WITH);
             DISPATCH();
         }
 

@@ -1003,11 +1003,6 @@ stack_effect(int opcode, int oparg, int jump)
         case INPLACE_OR:
             return -1;
 
-        case SETUP_WITH:
-            /* 1 in the normal flow.
-             * Restore the stack position and push 4 values before jumping to
-             * the handler if an exception is raised. */
-            return jump ? 4 : 1;
         case RETURN_VALUE:
             return -1;
         case IMPORT_STAR:
@@ -1152,7 +1147,9 @@ stack_effect(int opcode, int oparg, int jump)
              * of __aenter__ and push 7 values before jumping to the handler
              * if an exception be raised. */
             return jump ? -1 + 4 : 0;
+
         case BEFORE_ASYNC_WITH:
+        case BEFORE_WITH:
             return 1;
         case GET_AITER:
             return 0;
@@ -1793,6 +1790,7 @@ compiler_unwind_fblock(struct compiler *c, struct fblockinfo *info,
             if (preserve_tos) {
                 ADDOP(c, ROT_FOUR);
             }
+            ADDOP(c, POP_BLOCK);
             ADDOP(c, POP_EXCEPT);
             return 1;
 
@@ -1820,6 +1818,7 @@ compiler_unwind_fblock(struct compiler *c, struct fblockinfo *info,
             if (preserve_tos) {
                 ADDOP(c, ROT_FOUR);
             }
+            ADDOP(c, POP_BLOCK);
             ADDOP(c, POP_EXCEPT);
             if (info->fb_datum) {
                 ADDOP_LOAD_CONST(c, Py_None);
@@ -3188,6 +3187,7 @@ compiler_try_except(struct compiler *c, stmt_ty s)
             VISIT_SEQ(c, stmt, handler->v.ExceptHandler.body);
             compiler_pop_fblock(c, HANDLER_CLEANUP, cleanup_body);
             ADDOP(c, POP_BLOCK);
+            ADDOP(c, POP_BLOCK);
             ADDOP(c, POP_EXCEPT);
             /* name = None; del name; # Mark as artificial */
             c->u->u_lineno = -1;
@@ -3224,6 +3224,7 @@ compiler_try_except(struct compiler *c, stmt_ty s)
             compiler_pop_fblock(c, HANDLER_CLEANUP, cleanup_body);
             /* name = None; del name; # Mark as artificial */
             c->u->u_lineno = -1;
+            ADDOP(c, POP_BLOCK);
             ADDOP(c, POP_EXCEPT);
             ADDOP_JUMP(c, JUMP_FORWARD, end);
         }
@@ -4940,6 +4941,7 @@ compiler_with_except_finish(struct compiler *c, basicblock * cleanup) {
     ADDOP(c, POP_TOP);
     ADDOP(c, POP_TOP);
     ADDOP(c, POP_TOP);
+    ADDOP(c, POP_BLOCK);
     ADDOP(c, POP_EXCEPT);
     ADDOP(c, POP_TOP);
     ADDOP(c, POP_TOP);
@@ -5095,9 +5097,10 @@ compiler_with(struct compiler *c, stmt_ty s, int pos)
     /* Evaluate EXPR */
     VISIT(c, expr, item->context_expr);
     /* Will push bound __exit__ */
-    ADDOP_JUMP(c, SETUP_WITH, final);
+    ADDOP(c, BEFORE_WITH);
+    ADDOP_JUMP(c, SETUP_ASYNC_WITH, final);
 
-    /* SETUP_WITH pushes a finally block. */
+    /* SETUP_ASYNC_WITH pushes a finally block. */
     compiler_use_next_block(c, block);
     if (!compiler_push_fblock(c, WITH, block, final, NULL)) {
         return 0;
@@ -7416,7 +7419,6 @@ propogate_line_numbers(struct assembler *a) {
             switch (b->b_instr[b->b_iused-1].i_opcode) {
                 /* Note: Only actual jumps, not exception handlers */
                 case SETUP_ASYNC_WITH:
-                case SETUP_WITH:
                 case SETUP_FINALLY:
                 case SETUP_CLEANUP:
                     continue;
@@ -7523,7 +7525,6 @@ ensure_exits_have_lineno(struct compiler *c)
             switch (b->b_instr[b->b_iused-1].i_opcode) {
                 /* Note: Only actual jumps, not exception handlers */
                 case SETUP_ASYNC_WITH:
-                case SETUP_WITH:
                 case SETUP_FINALLY:
                 case SETUP_CLEANUP:
                     continue;