Produce cleaner bytecode for 'with' and 'async with' by generating separate code for normal and exceptional paths.

Author: markshannon

Include/opcode.h

@@ -30,10 +30,12 @@ extern "C" {
 #define BINARY_TRUE_DIVIDE       27
 #define INPLACE_FLOOR_DIVIDE     28
 #define INPLACE_TRUE_DIVIDE      29
+#define RERAISE                  48
 #define GET_AITER                50
 #define GET_ANEXT                51
 #define BEFORE_ASYNC_WITH        52
 #define BEGIN_FINALLY            53
+#define WITH_EXCEPT_START        54
 #define INPLACE_ADD              55
 #define INPLACE_SUBTRACT         56
 #define INPLACE_MULTIPLY         57

Lib/opcode.py

@@ -84,10 +84,13 @@ def jabs_op(name, op):
 def_op('INPLACE_FLOOR_DIVIDE', 28)
 def_op('INPLACE_TRUE_DIVIDE', 29)
 
+def_op('RERAISE', 48)
+
 def_op('GET_AITER', 50)
 def_op('GET_ANEXT', 51)
 def_op('BEFORE_ASYNC_WITH', 52)
 def_op('BEGIN_FINALLY', 53)
+def_op('WITH_EXCEPT_START', 54)
 
 def_op('INPLACE_ADD', 55)
 def_op('INPLACE_SUBTRACT', 56)

Objects/frameobject.c

@@ -264,7 +264,7 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
         PyTryBlock *b = &f->f_blockstack[f->f_iblock];
         delta = (f->f_stacktop - f->f_valuestack) - b->b_level;
         if (b->b_type == SETUP_FINALLY &&
-            code[b->b_handler] == WITH_CLEANUP_START)
+            code[b->b_handler] == WITH_EXCEPT_START)
         {
             /* Pop the exit function. */
             delta++;

Python/ceval.c

@@ -1844,6 +1844,15 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
+        TARGET(RERAISE) {
+            PyObject *exc = POP();
+            PyObject *val = POP();
+            PyObject *tb = POP();
+            assert(PyExceptionClass_Check(exc));
+            PyErr_Restore(exc, val, tb);
+            goto exception_unwind;
+        }
+
         TARGET(POP_FINALLY) {
             /* If oparg is 0 at the top of the stack are 1 or 6 values:
                Either:
@@ -2954,6 +2963,37 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
+        TARGET(WITH_EXCEPT_START) {
+            /* At the top of the stack are 7 values:
+               - (TOP, SECOND, THIRD) = exc_info()
+               - (FOURTH, FITH, SIXTH) = previous exception for EXCEPT_HANDLER
+               - SEVENTH: the context.__exit__ bound method
+               We call SEVENTH(TOP, SECOND, THIRD).
+               Then we push again the TOP exception and the __exit__
+               return value.
+            */
+            PyObject *exit_func;
+            PyObject *exit_stack[3];
+            PyObject *exc, *val, *tb, *res;
+
+            exc = TOP();
+            val = SECOND();
+            tb = THIRD();
+            assert(exc != Py_None);
+            assert(!PyLong_Check(exc));
+
+            exit_func = PEEK(7);
+            exit_stack[0] = exc;
+            exit_stack[1] = val;
+            exit_stack[2] = tb;
+            res = _PyObject_FastCall(exit_func, exit_stack, 3);
+            if (res == NULL)
+                goto error;
+
+            PUSH(res);
+            DISPATCH();
+        }
+
         TARGET(WITH_CLEANUP_START) {
             /* At the top of the stack are 1 or 6 values indicating
                how/why we entered the finally clause:

Python/compile.c

@@ -1052,9 +1052,14 @@ stack_effect(int opcode, int oparg, int jump)
              * This is the main reason of using this opcode instead of
              * "LOAD_CONST None". */
             return 6;
+        case RERAISE:
+            return -3;
         case CALL_FINALLY:
             return jump ? 1 : 0;
 
+        case WITH_EXCEPT_START:
+            return 1;
+
         case LOAD_FAST:
             return 1;
         case STORE_FAST:
@@ -1490,6 +1495,15 @@ compiler_pop_fblock(struct compiler *c, enum fblocktype t, basicblock *b)
     assert(u->u_fblock[u->u_nfblocks].fb_block == b);
 }
 
+static int
+compiler_call_exit_with_nones(struct compiler *c) {
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, DUP_TOP);
+    ADDOP(c, DUP_TOP);
+    ADDOP_I(c, CALL_FUNCTION, 3);
+    return 1;
+}
+
 /* Unwind a frame block.  If preserve_tos is true, the TOS before
  * popping the blocks will be restored afterwards.
  */
@@ -1528,15 +1542,15 @@ compiler_unwind_fblock(struct compiler *c, struct fblockinfo *info,
             if (preserve_tos) {
                 ADDOP(c, ROT_TWO);
             }
-            ADDOP(c, BEGIN_FINALLY);
-            ADDOP(c, WITH_CLEANUP_START);
+            if(!compiler_call_exit_with_nones(c)) {
+                return 0;
+            }
             if (info->fb_type == ASYNC_WITH) {
                 ADDOP(c, GET_AWAITABLE);
                 ADDOP_O(c, LOAD_CONST, Py_None, consts);
                 ADDOP(c, YIELD_FROM);
             }
-            ADDOP(c, WITH_CLEANUP_FINISH);
-            ADDOP_I(c, POP_FINALLY, 0);
+            ADDOP(c, POP_TOP);
             return 1;
 
         case HANDLER_CLEANUP:
@@ -4310,6 +4324,22 @@ expr_constant(expr_ty e)
     return -1;
 }
 
+static int
+compiler_with_except_finish(struct compiler *c) {
+    basicblock *exit;
+    exit = compiler_new_block(c);
+    if (exit == NULL)
+        return 0;
+    ADDOP_JABS(c, POP_JUMP_IF_TRUE, exit);
+    ADDOP(c, RERAISE);
+    compiler_use_next_block(c, exit);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_EXCEPT);
+    ADDOP(c, POP_TOP);
+    return 1;
+}
 
 /*
    Implements the async with statement.
@@ -4338,14 +4368,16 @@ expr_constant(expr_ty e)
 static int
 compiler_async_with(struct compiler *c, stmt_ty s, int pos)
 {
-    basicblock *block, *finally;
+    basicblock *block, *final1, *final2, *exit;
     withitem_ty item = asdl_seq_GET(s->v.AsyncWith.items, pos);
 
     assert(s->kind == AsyncWith_kind);
 
     block = compiler_new_block(c);
-    finally = compiler_new_block(c);
-    if (!block || !finally)
+    final1 = compiler_new_block(c);
+    final2 = compiler_new_block(c);
+    exit = compiler_new_block(c);
+    if (!block || !final1 || !final2 || !exit)
         return 0;
 
     /* Evaluate EXPR */
@@ -4356,11 +4388,11 @@ compiler_async_with(struct compiler *c, stmt_ty s, int pos)
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
 
-    ADDOP_JREL(c, SETUP_ASYNC_WITH, finally);
+    ADDOP_JREL(c, SETUP_ASYNC_WITH, final2);
 
     /* SETUP_ASYNC_WITH pushes a finally block. */
     compiler_use_next_block(c, block);
-    if (!compiler_push_fblock(c, ASYNC_WITH, block, finally)) {
+    if (!compiler_push_fblock(c, ASYNC_WITH, block, final2)) {
         return 0;
     }
 
@@ -4381,74 +4413,88 @@ compiler_async_with(struct compiler *c, stmt_ty s, int pos)
 
     /* End of try block; start the finally block */
     ADDOP(c, POP_BLOCK);
-    ADDOP(c, BEGIN_FINALLY);
     compiler_pop_fblock(c, ASYNC_WITH, block);
 
-    compiler_use_next_block(c, finally);
-    if (!compiler_push_fblock(c, FINALLY_END, finally, NULL))
+    compiler_use_next_block(c, final1);
+    if (!compiler_push_fblock(c, FINALLY_END, final1, NULL))
         return 0;
 
-    /* Finally block starts; context.__exit__ is on the stack under
-       the exception or return information. Just issue our magic
-       opcode. */
-    ADDOP(c, WITH_CLEANUP_START);
+    /* `finally` block for successful outcome:
+     * call __exit__(None, None, None)
+     */
+    if(!compiler_call_exit_with_nones(c))
+        return 0;
+    ADDOP(c, GET_AWAITABLE);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+
+    ADDOP(c, POP_TOP);
+
+    compiler_pop_fblock(c, FINALLY_END, final1);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, exit);
+
+    /* `finally` block for exceptional outcome */
+    compiler_use_next_block(c, final2);
+    if (!compiler_push_fblock(c, FINALLY_END, final2, NULL))
+        return 0;
 
+    ADDOP(c, WITH_EXCEPT_START);
     ADDOP(c, GET_AWAITABLE);
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
+    compiler_with_except_finish(c);
 
-    ADDOP(c, WITH_CLEANUP_FINISH);
+    compiler_pop_fblock(c, FINALLY_END, final2);
 
-    /* Finally block ends. */
-    ADDOP(c, END_FINALLY);
-    compiler_pop_fblock(c, FINALLY_END, finally);
+compiler_use_next_block(c, exit);
     return 1;
 }
 
 
 /*
    Implements the with statement from PEP 343.
-
-   The semantics outlined in that PEP are as follows:
-
    with EXPR as VAR:
        BLOCK
-
-   It is implemented roughly as:
-
-   context = EXPR
-   exit = context.__exit__  # not calling it
-   value = context.__enter__()
-   try:
-       VAR = value  # if VAR present in the syntax
-       BLOCK
-   finally:
-       if an exception was raised:
-           exc = copy of (exception, instance, traceback)
-       else:
-           exc = (None, None, None)
-       exit(*exc)
+   is implemented as:
+        <code for EXPR>
+        SETUP_WITH  E
+        <code to store to VAR> or POP_TOP
+        <code for BLOCK>
+        LOAD_CONST (None, None, None)
+        CALL_FUNCTION_EX 0
+        JUMP_FORWARD  EXIT
+    E:  WITH_EXCEPT_START (calls EXPR.__exit__)
+        POP_JUMP_IF_TRUE T:
+        RERAISE
+    T:  POP_TOP * 3 (remove exception from stack)
+        POP_EXCEPT
+        POP_TOP
+    EXIT:
  */
+
 static int
 compiler_with(struct compiler *c, stmt_ty s, int pos)
 {
-    basicblock *block, *finally;
+    basicblock *block, *final1, *final2, *exit;
     withitem_ty item = asdl_seq_GET(s->v.With.items, pos);
 
     assert(s->kind == With_kind);
 
     block = compiler_new_block(c);
-    finally = compiler_new_block(c);
-    if (!block || !finally)
+    final1 = compiler_new_block(c);
+    final2 = compiler_new_block(c);
+    exit = compiler_new_block(c);
+    if (!block || !final1 || !final2 || !exit)
         return 0;
 
     /* Evaluate EXPR */
     VISIT(c, expr, item->context_expr);
-    ADDOP_JREL(c, SETUP_WITH, finally);
+    /* Will push bound __exit__ */
+    ADDOP_JREL(c, SETUP_WITH, final2);
 
-    /* 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, finally)) {
+    if (!compiler_push_fblock(c, WITH, block, final2)) {
         return 0;
     }
 
@@ -4467,24 +4513,33 @@ compiler_with(struct compiler *c, stmt_ty s, int pos)
     else if (!compiler_with(c, s, pos))
             return 0;
 
-    /* End of try block; start the finally block */
+    /* End of try block; start the finally blocks */
     ADDOP(c, POP_BLOCK);
-    ADDOP(c, BEGIN_FINALLY);
     compiler_pop_fblock(c, WITH, block);
 
-    compiler_use_next_block(c, finally);
-    if (!compiler_push_fblock(c, FINALLY_END, finally, NULL))
+    /* `finally` block for successful outcome:
+     * call __exit__(None, None, None)
+     */
+    compiler_use_next_block(c, final1);
+    if (!compiler_push_fblock(c, FINALLY_END, final1, NULL))
         return 0;
 
-    /* Finally block starts; context.__exit__ is on the stack under
-       the exception or return information. Just issue our magic
-       opcode. */
-    ADDOP(c, WITH_CLEANUP_START);
-    ADDOP(c, WITH_CLEANUP_FINISH);
+    if(!compiler_call_exit_with_nones(c))
+        return 0;
+    ADDOP(c, POP_TOP);
+    compiler_pop_fblock(c, FINALLY_END, final1);
+    ADDOP_JREL(c, JUMP_FORWARD, exit);
 
-    /* Finally block ends. */
-    ADDOP(c, END_FINALLY);
-    compiler_pop_fblock(c, FINALLY_END, finally);
+    /* `finally` block for exceptional outcome */
+    compiler_use_next_block(c, final2);
+    if (!compiler_push_fblock(c, FINALLY_END, final2, NULL))
+        return 0;
+
+    ADDOP(c, WITH_EXCEPT_START);
+    compiler_with_except_finish(c);
+    compiler_pop_fblock(c, FINALLY_END, final2);
+
+    compiler_use_next_block(c, exit);
     return 1;
 }
 

Python/opcode_targets.h

@@ -47,13 +47,13 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
+    &&TARGET_RERAISE,
     &&_unknown_opcode,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
     &&TARGET_BEFORE_ASYNC_WITH,
     &&TARGET_BEGIN_FINALLY,
-    &&_unknown_opcode,
+    &&TARGET_WITH_EXCEPT_START,
     &&TARGET_INPLACE_ADD,
     &&TARGET_INPLACE_SUBTRACT,
     &&TARGET_INPLACE_MULTIPLY,