bpo-17611. Move unwinding of stack for "pseudo exceptions" from interpreter to compiler. (GH-5006)

Co-authored-by: Mark Shannon <mark@hotpy.org>
Co-authored-by: Antoine Pitrou <antoine@python.org>

2 files changed, 81 insertions, 132 deletions

diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 3083e5b..d308457 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -45,6 +45,27 @@ frame_getlineno(PyFrameObject *f, void *closure)
     return PyLong_FromLong(PyFrame_GetLineNumber(f));
 }
 
+
+/* Given the index of the effective opcode,
+   scan back to construct the oparg with EXTENDED_ARG */
+static unsigned int
+get_arg(const _Py_CODEUNIT *codestr, Py_ssize_t i)
+{
+    _Py_CODEUNIT word;
+    unsigned int oparg = _Py_OPARG(codestr[i]);
+    if (i >= 1 && _Py_OPCODE(word = codestr[i-1]) == EXTENDED_ARG) {
+        oparg |= _Py_OPARG(word) << 8;
+        if (i >= 2 && _Py_OPCODE(word = codestr[i-2]) == EXTENDED_ARG) {
+            oparg |= _Py_OPARG(word) << 16;
+            if (i >= 3 && _Py_OPCODE(word = codestr[i-3]) == EXTENDED_ARG) {
+                oparg |= _Py_OPARG(word) << 24;
+            }
+        }
+    }
+    return oparg;
+}
+
+
 /* Setter for f_lineno - you can set f_lineno from within a trace function in
  * order to jump to a given line of code, subject to some restrictions.  Most
  * lines are OK to jump to because they don't make any assumptions about the
@@ -56,8 +77,9 @@ frame_getlineno(PyFrameObject *f, void *closure)
  *    they expect an exception to be on the top of the stack.
  *  o Lines that live in a 'finally' block can't be jumped from or to, since
  *    the END_FINALLY expects to clean up the stack after the 'try' block.
- *  o 'try'/'for'/'while' blocks can't be jumped into because the blockstack
- *    needs to be set up before their code runs, and for 'for' loops the
+ *  o 'try', 'with' and 'async with' blocks can't be jumped into because
+ *    the blockstack needs to be set up before their code runs.
+ *  o 'for' and 'async for' loops can't be jumped into because the
  *    iterator needs to be on the stack.
  */
 static int
@@ -75,17 +97,10 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
     int offset = 0;                     /* (ditto) */
     int line = 0;                       /* (ditto) */
     int addr = 0;                       /* (ditto) */
-    int min_addr = 0;                   /* Scanning the SETUPs and POPs */
-    int max_addr = 0;                   /* (ditto) */
-    int delta_iblock = 0;               /* (ditto) */
-    int min_delta_iblock = 0;           /* (ditto) */
-    int min_iblock = 0;                 /* (ditto) */
-    int f_lasti_setup_addr = 0;         /* Policing no-jump-into-finally */
-    int new_lasti_setup_addr = 0;       /* (ditto) */
+    int delta_iblock = 0;               /* Scanning the SETUPs and POPs */
+    int for_loop_delta = 0;             /* (ditto) */
     int blockstack[CO_MAXBLOCKS];       /* Walking the 'finally' blocks */
-    int in_finally[CO_MAXBLOCKS];       /* (ditto) */
     int blockstack_top = 0;             /* (ditto) */
-    unsigned char setup_op = 0;         /* (ditto) */
 
     /* f_lineno must be an integer. */
     if (!PyLong_CheckExact(p_new_lineno)) {
@@ -159,8 +174,6 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
 
     /* We're now ready to look at the bytecode. */
     PyBytes_AsStringAndSize(f->f_code->co_code, (char **)&code, &code_len);
-    min_addr = Py_MIN(new_lasti, f->f_lasti);
-    max_addr = Py_MAX(new_lasti, f->f_lasti);
 
     /* You can't jump onto a line with an 'except' statement on it -
      * they expect to have an exception on the top of the stack, which
@@ -179,141 +192,73 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
     }
 
     /* You can't jump into or out of a 'finally' block because the 'try'
-     * block leaves something on the stack for the END_FINALLY to clean
-     * up.      So we walk the bytecode, maintaining a simulated blockstack.
-     * When we reach the old or new address and it's in a 'finally' block
-     * we note the address of the corresponding SETUP_FINALLY.  The jump
-     * is only legal if neither address is in a 'finally' block or
-     * they're both in the same one.  'blockstack' is a stack of the
-     * bytecode addresses of the SETUP_X opcodes, and 'in_finally' tracks
-     * whether we're in a 'finally' block at each blockstack level. */
-    f_lasti_setup_addr = -1;
-    new_lasti_setup_addr = -1;
+     * block leaves something on the stack for the END_FINALLY to clean up.
+     * So we walk the bytecode, maintaining a simulated blockstack.
+     * 'blockstack' is a stack of the bytecode addresses of the starts of
+     * the 'finally' blocks. */
     memset(blockstack, '\0', sizeof(blockstack));
-    memset(in_finally, '\0', sizeof(in_finally));
     blockstack_top = 0;
     for (addr = 0; addr < code_len; addr += sizeof(_Py_CODEUNIT)) {
         unsigned char op = code[addr];
         switch (op) {
-        case SETUP_LOOP:
-        case SETUP_EXCEPT:
         case SETUP_FINALLY:
         case SETUP_WITH:
         case SETUP_ASYNC_WITH:
-            blockstack[blockstack_top++] = addr;
-            in_finally[blockstack_top-1] = 0;
-            break;
-
-        case POP_BLOCK:
-            assert(blockstack_top > 0);
-            setup_op = code[blockstack[blockstack_top-1]];
-            if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH
-                                    || setup_op == SETUP_ASYNC_WITH) {
-                in_finally[blockstack_top-1] = 1;
-            }
-            else {
-                blockstack_top--;
+        case FOR_ITER: {
+            unsigned int oparg = get_arg((const _Py_CODEUNIT *)code,
+                                         addr / sizeof(_Py_CODEUNIT));
+            int target_addr = addr + oparg + sizeof(_Py_CODEUNIT);
+            assert(target_addr < code_len);
+            /* Police block-jumping (you can't jump into the middle of a block)
+             * and ensure that the blockstack finishes up in a sensible state (by
+             * popping any blocks we're jumping out of).  We look at all the
+             * blockstack operations between the current position and the new
+             * one, and keep track of how many blocks we drop out of on the way.
+             * By also keeping track of the lowest blockstack position we see, we
+             * can tell whether the jump goes into any blocks without coming out
+             * again - in that case we raise an exception below. */
+            int first_in = addr < f->f_lasti && f->f_lasti < target_addr;
+            int second_in = addr < new_lasti && new_lasti < target_addr;
+            if (!first_in && second_in) {
+                PyErr_SetString(PyExc_ValueError,
+                                "can't jump into the middle of a block");
+                return -1;
             }
-            break;
-
-        case END_FINALLY:
-            /* Ignore END_FINALLYs for SETUP_EXCEPTs - they exist
-             * in the bytecode but don't correspond to an actual
-             * 'finally' block.  (If blockstack_top is 0, we must
-             * be seeing such an END_FINALLY.) */
-            if (blockstack_top > 0) {
-                setup_op = code[blockstack[blockstack_top-1]];
-                if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH
-                                    || setup_op == SETUP_ASYNC_WITH) {
-                    blockstack_top--;
+            if (first_in && !second_in) {
+                if (op == FOR_ITER && !delta_iblock) {
+                    for_loop_delta++;
+                }
+                if (op != FOR_ITER) {
+                    delta_iblock++;
                 }
             }
+            if (op != FOR_ITER) {
+                blockstack[blockstack_top++] = target_addr;
+            }
             break;
         }
 
-        /* For the addresses we're interested in, see whether they're
-         * within a 'finally' block and if so, remember the address
-         * of the SETUP_FINALLY. */
-        if (addr == new_lasti || addr == f->f_lasti) {
-            int i = 0;
-            int setup_addr = -1;
-            for (i = blockstack_top-1; i >= 0; i--) {
-                if (in_finally[i]) {
-                    setup_addr = blockstack[i];
-                    break;
-                }
-            }
-
-            if (setup_addr != -1) {
-                if (addr == new_lasti) {
-                    new_lasti_setup_addr = setup_addr;
-                }
-
-                if (addr == f->f_lasti) {
-                    f_lasti_setup_addr = setup_addr;
-                }
+        case END_FINALLY: {
+            assert(blockstack_top > 0);
+            int target_addr = blockstack[--blockstack_top];
+            assert(target_addr <= addr);
+            int first_in = target_addr <= f->f_lasti && f->f_lasti <= addr;
+            int second_in = target_addr <= new_lasti && new_lasti <= addr;
+            if (first_in != second_in) {
+                PyErr_SetString(PyExc_ValueError,
+                                "can't jump into or out of a 'finally' block");
+                return -1;
             }
+            break;
+        }
         }
     }
 
     /* Verify that the blockstack tracking code didn't get lost. */
     assert(blockstack_top == 0);
 
-    /* After all that, are we jumping into / out of a 'finally' block? */
-    if (new_lasti_setup_addr != f_lasti_setup_addr) {
-        PyErr_SetString(PyExc_ValueError,
-                    "can't jump into or out of a 'finally' block");
-        return -1;
-    }
-
-
-    /* Police block-jumping (you can't jump into the middle of a block)
-     * and ensure that the blockstack finishes up in a sensible state (by
-     * popping any blocks we're jumping out of).  We look at all the
-     * blockstack operations between the current position and the new
-     * one, and keep track of how many blocks we drop out of on the way.
-     * By also keeping track of the lowest blockstack position we see, we
-     * can tell whether the jump goes into any blocks without coming out
-     * again - in that case we raise an exception below. */
-    delta_iblock = 0;
-    for (addr = min_addr; addr < max_addr; addr += sizeof(_Py_CODEUNIT)) {
-        unsigned char op = code[addr];
-        switch (op) {
-        case SETUP_LOOP:
-        case SETUP_EXCEPT:
-        case SETUP_FINALLY:
-        case SETUP_WITH:
-        case SETUP_ASYNC_WITH:
-            delta_iblock++;
-            break;
-
-        case POP_BLOCK:
-            delta_iblock--;
-            break;
-        }
-
-        min_delta_iblock = Py_MIN(min_delta_iblock, delta_iblock);
-    }
-
-    /* Derive the absolute iblock values from the deltas. */
-    min_iblock = f->f_iblock + min_delta_iblock;
-    if (new_lasti > f->f_lasti) {
-        /* Forwards jump. */
-        new_iblock = f->f_iblock + delta_iblock;
-    }
-    else {
-        /* Backwards jump. */
-        new_iblock = f->f_iblock - delta_iblock;
-    }
-
-    /* Are we jumping into a block? */
-    if (new_iblock > min_iblock) {
-        PyErr_SetString(PyExc_ValueError,
-                        "can't jump into the middle of a block");
-        return -1;
-    }
-
     /* Pop any blocks that we're jumping out of. */
+    new_iblock = f->f_iblock - delta_iblock;
     while (f->f_iblock > new_iblock) {
         PyTryBlock *b = &f->f_blockstack[--f->f_iblock];
         while ((f->f_stacktop - f->f_valuestack) > b->b_level) {
@@ -321,6 +266,12 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
             Py_DECREF(v);
         }
     }
+    /* Pop the iterators of any 'for' loop we're jumping out of. */
+    while (for_loop_delta > 0) {
+        PyObject *v = (*--f->f_stacktop);
+        Py_DECREF(v);
+        for_loop_delta--;
+    }
 
     /* Finally set the new f_lineno and f_lasti and return OK. */
     f->f_lineno = new_lineno;
diff --git a/Objects/genobject.c b/Objects/genobject.c
index 88b03c5..9f59338 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -822,18 +822,16 @@ PyGen_New(PyFrameObject *f)
 int
 PyGen_NeedsFinalizing(PyGenObject *gen)
 {
-    int i;
     PyFrameObject *f = gen->gi_frame;
 
     if (f == NULL || f->f_stacktop == NULL)
         return 0; /* no frame or empty blockstack == no finalization */
 
-    /* Any block type besides a loop requires cleanup. */
-    for (i = 0; i < f->f_iblock; i++)
-        if (f->f_blockstack[i].b_type != SETUP_LOOP)
-            return 1;
+    /* Any (exception-handling) block type requires cleanup. */
+    if (f->f_iblock > 0)
+        return 1;
 
-    /* No blocks except loops, it's safe to skip finalization. */
+    /* No blocks, it's safe to skip finalization. */
     return 0;
 }