bpo-40222: "Zero cost" exception handling (GH-25729)

"Zero cost" exception handling.

* Uses a lookup table to determine how to handle exceptions.
* Removes SETUP_FINALLY and POP_TOP block instructions, eliminating (most of) the runtime overhead of try statements.
* Reduces the size of the frame object by about 60%.

4 files changed, 405 insertions, 197 deletions

diff --git a/Objects/clinic/codeobject.c.h b/Objects/clinic/codeobject.c.h
index bae2ab0..7ffdf07 100644
--- a/Objects/clinic/codeobject.c.h
+++ b/Objects/clinic/codeobject.c.h
@@ -5,7 +5,8 @@ preserve
 PyDoc_STRVAR(code_new__doc__,
 "code(argcount, posonlyargcount, kwonlyargcount, nlocals, stacksize,\n"
 "     flags, codestring, constants, names, varnames, filename, name,\n"
-"     firstlineno, linetable, freevars=(), cellvars=(), /)\n"
+"     firstlineno, linetable, exceptiontable, freevars=(), cellvars=(),\n"
+"     /)\n"
 "--\n"
 "\n"
 "Create a code object.  Not for the faint of heart.");
@@ -15,8 +16,8 @@ code_new_impl(PyTypeObject *type, int argcount, int posonlyargcount,
               int kwonlyargcount, int nlocals, int stacksize, int flags,
               PyObject *code, PyObject *consts, PyObject *names,
               PyObject *varnames, PyObject *filename, PyObject *name,
-              int firstlineno, PyObject *linetable, PyObject *freevars,
-              PyObject *cellvars);
+              int firstlineno, PyObject *linetable, PyObject *exceptiontable,
+              PyObject *freevars, PyObject *cellvars);
 
 static PyObject *
 code_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
@@ -36,6 +37,7 @@ code_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
     PyObject *name;
     int firstlineno;
     PyObject *linetable;
+    PyObject *exceptiontable;
     PyObject *freevars = NULL;
     PyObject *cellvars = NULL;
 
@@ -43,7 +45,7 @@ code_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
         !_PyArg_NoKeywords("code", kwargs)) {
         goto exit;
     }
-    if (!_PyArg_CheckPositional("code", PyTuple_GET_SIZE(args), 14, 16)) {
+    if (!_PyArg_CheckPositional("code", PyTuple_GET_SIZE(args), 15, 17)) {
         goto exit;
     }
     argcount = _PyLong_AsInt(PyTuple_GET_ITEM(args, 0));
@@ -115,14 +117,11 @@ code_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
         goto exit;
     }
     linetable = PyTuple_GET_ITEM(args, 13);
-    if (PyTuple_GET_SIZE(args) < 15) {
-        goto skip_optional;
-    }
-    if (!PyTuple_Check(PyTuple_GET_ITEM(args, 14))) {
-        _PyArg_BadArgument("code", "argument 15", "tuple", PyTuple_GET_ITEM(args, 14));
+    if (!PyBytes_Check(PyTuple_GET_ITEM(args, 14))) {
+        _PyArg_BadArgument("code", "argument 15", "bytes", PyTuple_GET_ITEM(args, 14));
         goto exit;
     }
-    freevars = PyTuple_GET_ITEM(args, 14);
+    exceptiontable = PyTuple_GET_ITEM(args, 14);
     if (PyTuple_GET_SIZE(args) < 16) {
         goto skip_optional;
     }
@@ -130,9 +129,17 @@ code_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
         _PyArg_BadArgument("code", "argument 16", "tuple", PyTuple_GET_ITEM(args, 15));
         goto exit;
     }
-    cellvars = PyTuple_GET_ITEM(args, 15);
+    freevars = PyTuple_GET_ITEM(args, 15);
+    if (PyTuple_GET_SIZE(args) < 17) {
+        goto skip_optional;
+    }
+    if (!PyTuple_Check(PyTuple_GET_ITEM(args, 16))) {
+        _PyArg_BadArgument("code", "argument 17", "tuple", PyTuple_GET_ITEM(args, 16));
+        goto exit;
+    }
+    cellvars = PyTuple_GET_ITEM(args, 16);
 skip_optional:
-    return_value = code_new_impl(type, argcount, posonlyargcount, kwonlyargcount, nlocals, stacksize, flags, code, consts, names, varnames, filename, name, firstlineno, linetable, freevars, cellvars);
+    return_value = code_new_impl(type, argcount, posonlyargcount, kwonlyargcount, nlocals, stacksize, flags, code, consts, names, varnames, filename, name, firstlineno, linetable, exceptiontable, freevars, cellvars);
 
 exit:
     return return_value;
@@ -144,7 +151,7 @@ PyDoc_STRVAR(code_replace__doc__,
 "        co_flags=-1, co_firstlineno=-1, co_code=None, co_consts=None,\n"
 "        co_names=None, co_varnames=None, co_freevars=None,\n"
 "        co_cellvars=None, co_filename=None, co_name=None,\n"
-"        co_linetable=None)\n"
+"        co_linetable=None, co_exceptiontable=None)\n"
 "--\n"
 "\n"
 "Return a copy of the code object with new values for the specified fields.");
@@ -160,15 +167,16 @@ code_replace_impl(PyCodeObject *self, int co_argcount,
                   PyObject *co_consts, PyObject *co_names,
                   PyObject *co_varnames, PyObject *co_freevars,
                   PyObject *co_cellvars, PyObject *co_filename,
-                  PyObject *co_name, PyBytesObject *co_linetable);
+                  PyObject *co_name, PyBytesObject *co_linetable,
+                  PyBytesObject *co_exceptiontable);
 
 static PyObject *
 code_replace(PyCodeObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
 {
     PyObject *return_value = NULL;
-    static const char * const _keywords[] = {"co_argcount", "co_posonlyargcount", "co_kwonlyargcount", "co_nlocals", "co_stacksize", "co_flags", "co_firstlineno", "co_code", "co_consts", "co_names", "co_varnames", "co_freevars", "co_cellvars", "co_filename", "co_name", "co_linetable", NULL};
+    static const char * const _keywords[] = {"co_argcount", "co_posonlyargcount", "co_kwonlyargcount", "co_nlocals", "co_stacksize", "co_flags", "co_firstlineno", "co_code", "co_consts", "co_names", "co_varnames", "co_freevars", "co_cellvars", "co_filename", "co_name", "co_linetable", "co_exceptiontable", NULL};
     static _PyArg_Parser _parser = {NULL, _keywords, "replace", 0};
-    PyObject *argsbuf[16];
+    PyObject *argsbuf[17];
     Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0;
     int co_argcount = self->co_argcount;
     int co_posonlyargcount = self->co_posonlyargcount;
@@ -186,6 +194,7 @@ code_replace(PyCodeObject *self, PyObject *const *args, Py_ssize_t nargs, PyObje
     PyObject *co_filename = self->co_filename;
     PyObject *co_name = self->co_name;
     PyBytesObject *co_linetable = (PyBytesObject *)self->co_linetable;
+    PyBytesObject *co_exceptiontable = (PyBytesObject *)self->co_exceptiontable;
 
     args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 0, 0, argsbuf);
     if (!args) {
@@ -343,15 +352,25 @@ code_replace(PyCodeObject *self, PyObject *const *args, Py_ssize_t nargs, PyObje
             goto skip_optional_kwonly;
         }
     }
-    if (!PyBytes_Check(args[15])) {
-        _PyArg_BadArgument("replace", "argument 'co_linetable'", "bytes", args[15]);
+    if (args[15]) {
+        if (!PyBytes_Check(args[15])) {
+            _PyArg_BadArgument("replace", "argument 'co_linetable'", "bytes", args[15]);
+            goto exit;
+        }
+        co_linetable = (PyBytesObject *)args[15];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (!PyBytes_Check(args[16])) {
+        _PyArg_BadArgument("replace", "argument 'co_exceptiontable'", "bytes", args[16]);
         goto exit;
     }
-    co_linetable = (PyBytesObject *)args[15];
+    co_exceptiontable = (PyBytesObject *)args[16];
 skip_optional_kwonly:
-    return_value = code_replace_impl(self, co_argcount, co_posonlyargcount, co_kwonlyargcount, co_nlocals, co_stacksize, co_flags, co_firstlineno, co_code, co_consts, co_names, co_varnames, co_freevars, co_cellvars, co_filename, co_name, co_linetable);
+    return_value = code_replace_impl(self, co_argcount, co_posonlyargcount, co_kwonlyargcount, co_nlocals, co_stacksize, co_flags, co_firstlineno, co_code, co_consts, co_names, co_varnames, co_freevars, co_cellvars, co_filename, co_name, co_linetable, co_exceptiontable);
 
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=e3091c7baaaaa420 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=a272b22f63ea002e input=a9049054013a1b77]*/
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index c76ac90..e981e39 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -119,7 +119,7 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
                           PyObject *code, PyObject *consts, PyObject *names,
                           PyObject *varnames, PyObject *freevars, PyObject *cellvars,
                           PyObject *filename, PyObject *name, int firstlineno,
-                          PyObject *linetable)
+                          PyObject *linetable, PyObject *exceptiontable)
 {
     PyCodeObject *co;
     Py_ssize_t *cell2arg = NULL;
@@ -137,7 +137,8 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
         cellvars == NULL || !PyTuple_Check(cellvars) ||
         name == NULL || !PyUnicode_Check(name) ||
         filename == NULL || !PyUnicode_Check(filename) ||
-        linetable == NULL || !PyBytes_Check(linetable)) {
+        linetable == NULL || !PyBytes_Check(linetable) ||
+        exceptiontable == NULL || !PyBytes_Check(exceptiontable)) {
         PyErr_BadInternalCall();
         return NULL;
     }
@@ -260,6 +261,8 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
     co->co_firstlineno = firstlineno;
     Py_INCREF(linetable);
     co->co_linetable = linetable;
+    Py_INCREF(exceptiontable);
+    co->co_exceptiontable = exceptiontable;
     co->co_zombieframe = NULL;
     co->co_weakreflist = NULL;
     co->co_extra = NULL;
@@ -277,12 +280,12 @@ PyCode_New(int argcount, int kwonlyargcount,
            PyObject *code, PyObject *consts, PyObject *names,
            PyObject *varnames, PyObject *freevars, PyObject *cellvars,
            PyObject *filename, PyObject *name, int firstlineno,
-           PyObject *linetable)
+           PyObject *linetable, PyObject *exceptiontable)
 {
     return PyCode_NewWithPosOnlyArgs(argcount, 0, kwonlyargcount, nlocals,
                                      stacksize, flags, code, consts, names,
                                      varnames, freevars, cellvars, filename,
-                                     name, firstlineno, linetable);
+                                     name, firstlineno, linetable, exceptiontable);
 }
 
 int
@@ -369,7 +372,8 @@ PyCode_NewEmpty(const char *filename, const char *funcname, int firstlineno)
                 filename_ob,                    /* filename */
                 funcname_ob,                    /* name */
                 firstlineno,                    /* firstlineno */
-                emptystring                     /* linetable */
+                emptystring,                    /* linetable */
+                emptystring                     /* exception table */
                 );
 
 failed:
@@ -397,6 +401,7 @@ static PyMemberDef code_memberlist[] = {
     {"co_name",         T_OBJECT,       OFF(co_name),            READONLY},
     {"co_firstlineno",  T_INT,          OFF(co_firstlineno),     READONLY},
     {"co_linetable",    T_OBJECT,       OFF(co_linetable),       READONLY},
+    {"co_exceptiontable",    T_OBJECT,  OFF(co_exceptiontable),  READONLY},
     {NULL}      /* Sentinel */
 };
 
@@ -538,6 +543,7 @@ code.__new__ as code_new
     name: unicode
     firstlineno: int
     linetable: object(subclass_of="&PyBytes_Type")
+    exceptiontable: object(subclass_of="&PyBytes_Type")
     freevars: object(subclass_of="&PyTuple_Type", c_default="NULL") = ()
     cellvars: object(subclass_of="&PyTuple_Type", c_default="NULL") = ()
     /
@@ -550,9 +556,9 @@ code_new_impl(PyTypeObject *type, int argcount, int posonlyargcount,
               int kwonlyargcount, int nlocals, int stacksize, int flags,
               PyObject *code, PyObject *consts, PyObject *names,
               PyObject *varnames, PyObject *filename, PyObject *name,
-              int firstlineno, PyObject *linetable, PyObject *freevars,
-              PyObject *cellvars)
-/*[clinic end generated code: output=42c1839b082ba293 input=0ec80da632b99f57]*/
+              int firstlineno, PyObject *linetable, PyObject *exceptiontable,
+              PyObject *freevars, PyObject *cellvars)
+/*[clinic end generated code: output=a3899259c3b4cace input=f823c686da4b3a03]*/
 {
     PyObject *co = NULL;
     PyObject *ournames = NULL;
@@ -618,7 +624,9 @@ code_new_impl(PyTypeObject *type, int argcount, int posonlyargcount,
                                                code, consts, ournames,
                                                ourvarnames, ourfreevars,
                                                ourcellvars, filename,
-                                               name, firstlineno, linetable);
+                                               name, firstlineno, linetable,
+                                               exceptiontable
+                                              );
   cleanup:
     Py_XDECREF(ournames);
     Py_XDECREF(ourvarnames);
@@ -663,6 +671,7 @@ code_dealloc(PyCodeObject *co)
     Py_XDECREF(co->co_filename);
     Py_XDECREF(co->co_name);
     Py_XDECREF(co->co_linetable);
+    Py_XDECREF(co->co_exceptiontable);
     if (co->co_cell2arg != NULL)
         PyMem_Free(co->co_cell2arg);
     if (co->co_zombieframe != NULL)
@@ -715,6 +724,7 @@ code.replace
     co_filename: unicode(c_default="self->co_filename") = None
     co_name: unicode(c_default="self->co_name") = None
     co_linetable: PyBytesObject(c_default="(PyBytesObject *)self->co_linetable") = None
+    co_exceptiontable: PyBytesObject(c_default="(PyBytesObject *)self->co_exceptiontable") = None
 
 Return a copy of the code object with new values for the specified fields.
 [clinic start generated code]*/
@@ -727,8 +737,9 @@ code_replace_impl(PyCodeObject *self, int co_argcount,
                   PyObject *co_consts, PyObject *co_names,
                   PyObject *co_varnames, PyObject *co_freevars,
                   PyObject *co_cellvars, PyObject *co_filename,
-                  PyObject *co_name, PyBytesObject *co_linetable)
-/*[clinic end generated code: output=50d77e668d3b449b input=a5f997b173d7f636]*/
+                  PyObject *co_name, PyBytesObject *co_linetable,
+                  PyBytesObject *co_exceptiontable)
+/*[clinic end generated code: output=80957472b7f78ed6 input=38376b1193efbbae]*/
 {
 #define CHECK_INT_ARG(ARG) \
         if (ARG < 0) { \
@@ -758,7 +769,7 @@ code_replace_impl(PyCodeObject *self, int co_argcount,
         co_argcount, co_posonlyargcount, co_kwonlyargcount, co_nlocals,
         co_stacksize, co_flags, (PyObject*)co_code, co_consts, co_names,
         co_varnames, co_freevars, co_cellvars, co_filename, co_name,
-        co_firstlineno, (PyObject*)co_linetable);
+        co_firstlineno, (PyObject*)co_linetable, (PyObject*)co_exceptiontable);
 }
 
 static PyObject *
diff --git a/Objects/exception_handling_notes.txt b/Objects/exception_handling_notes.txt
new file mode 100644
index 0000000..2183fa1
--- /dev/null
+++ b/Objects/exception_handling_notes.txt
@@ -0,0 +1,179 @@
+Description of exception handling in Python 3.11
+------------------------------------------------
+
+Python 3.11 uses what is known as "zero-cost" exception handling.
+Prior to 3.11, exceptions were handled by a runtime stack of "blocks".
+
+In zero-cost exception handling, the cost of supporting exceptions is minimized.
+In the common case (where no exception is raised) the cost is reduced
+to zero (or close to zero).
+The cost of raising an exception is increased, but not by much.
+
+The following code:
+
+def f():
+    try:
+        g(0)
+    except:
+        return "fail"
+
+compiles as follows in 3.10:
+
+  2           0 SETUP_FINALLY            7 (to 16)
+
+  3           2 LOAD_GLOBAL              0 (g)
+              4 LOAD_CONST               1 (0)
+              6 CALL_FUNCTION            1
+              8 POP_TOP
+             10 POP_BLOCK
+             12 LOAD_CONST               0 (None)
+             14 RETURN_VALUE
+
+  4     >>   16 POP_TOP
+             18 POP_TOP
+             20 POP_TOP
+
+  5          22 POP_EXCEPT
+             24 LOAD_CONST               3 ('fail')
+             26 RETURN_VALUE
+
+Note the explicit instructions to push and pop from the "block" stack:
+SETUP_FINALLY and POP_BLOCK.
+
+In 3.11, the SETUP_FINALLY and POP_BLOCK are eliminated, replaced with
+a table to determine where to jump to when an exception is raised.
+
+  2           0 NOP
+
+  3           2 LOAD_GLOBAL              0 (g)
+              4 LOAD_CONST               1 (0)
+              6 CALL_FUNCTION            1
+              8 POP_TOP
+             10 LOAD_CONST               0 (None)
+             12 RETURN_VALUE
+        >>   14 PUSH_EXC_INFO
+
+  4          16 POP_TOP
+             18 POP_TOP
+             20 POP_TOP
+
+  5          22 POP_EXCEPT
+             24 LOAD_CONST               2 ('fail')
+             26 RETURN_VALUE
+        >>   28 POP_EXCEPT_AND_RERAISE
+ExceptionTable:
+  2 to 8 -> 14 [0]
+  14 to 20 -> 28 [3] lasti
+
+(Note this code is from an early 3.11 alpha, the NOP may well have be removed before release).
+
+If an instruction raises an exception then its offset is used to find the target to jump to.
+For example, the CALL_FUNCTION at offset 6, falls into the range 2 to 8.
+So, if g() raises an exception, then control jumps to offset 14.
+
+
+Unwinding
+---------
+
+When an exception is raised, the current instruction offset is used to find following:
+target to jump to, stack depth, and 'lasti', which determines whether the instruction
+offset of the raising instruction should be pushed.
+
+This information is stored in the exception table, described below.
+
+If there is no relevant entry, the exception bubbles up to the caller.
+
+If there is an entry, then:
+ 1. pop values from the stack until it matches the stack depth for the handler,
+ 2. if 'lasti' is true, then push the offset that the exception was raised at.
+ 3. push the exception to the stack as three values: traceback, value, type,
+ 4. jump to the target offset and resume execution.
+
+
+Format of the exception table
+-----------------------------
+
+Conceptually, the exception table consists of a sequence of 5-tuples:
+    1. start-offset (inclusive)
+    2. end-offset (exclusive)
+    3. target
+    4. stack-depth
+    5. push-lasti (boolean)
+
+All offsets and lengths are in instructions, not bytes.
+
+We want the format to be compact, but quickly searchable.
+For it to be compact, it needs to have variable sized entries so that we can store common (small) offsets compactly, but handle large offsets if needed.
+For it to be searchable quickly, we need to support binary search giving us log(n) performance in all cases.
+Binary search typically assumes fixed size entries, but that is not necesary, as long as we can identify the start of an entry.
+
+It is worth noting that the size (end-start) is always smaller than the end, so we encode the entries as:
+    start, size, target, depth, push-lasti
+
+Also, sizes are limited to 2**30 as the code length cannot exceed 2**31 and each instruction takes 2 bytes.
+It also happens that depth is generally quite small.
+
+So, we need to encode:
+    start (up to 30 bits)
+    size (up to 30 bits)
+    target (up to 30 bits)
+    depth (up to ~8 bits)
+    lasti (1 bit)
+
+We need a marker for the start of the entry, so the first byte of entry will have the most significant bit set.
+Since the most significant bit is reserved for marking the start of an entry, we have 7 bits per byte to encode offsets.
+Encoding uses a standard varint encoding, but with only 7 bits instead of the usual 8.
+The 8 bits of a bit are (msb left) SXdddddd where S is the start bit. X is the extend bit meaning that the next byte is required to extend the offset.
+
+In addition, we will combine depth and lasti into a single value, ((depth<<1)+lasti), before encoding.
+
+For example, the exception entry:
+    start:  20
+    end:    28
+    target: 100
+    depth:  3
+    lasti:  False
+
+is encoded first by converting to the more compact four value form:
+    start:         20
+    size:          8
+    target:        100
+  depth<<1+lasti:  6
+
+which is then encoded as:
+    148 (MSB + 20 for start)
+    8   (size)
+    65  (Extend bit + 1)
+    36  (Remainder of target, 100 == (1<<6)+36)
+    6
+
+for a total of five bytes.
+
+
+
+Script to parse the exception table
+-----------------------------------
+
+def parse_varint(iterator):
+    b = next(iterator)
+    val = b & 63
+    while b&64:
+        val <<= 6
+        b = next(iterator)
+        val |= b&63
+    return val
+
+def parse_exception_table(code):
+    iterator = iter(code.co_exceptiontable)
+    try:
+        while True:
+            start = parse_varint(iterator)*2
+            length = parse_varint(iterator)*2
+            end = start + length - 2 # Present as inclusive, not exclusive
+            target = parse_varint(iterator)*2
+            dl = parse_varint(iterator)
+            depth = dl >> 1
+            lasti = bool(dl&1)
+            yield start, end, target, depth, lasti
+    except StopIteration:
+        return
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 034b908..ae8cdcf 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -91,56 +91,71 @@ get_arg(const _Py_CODEUNIT *codestr, Py_ssize_t i)
     return oparg;
 }
 
+/* Model the evaluation stack, to determine which jumps
+ * are safe and how many values needs to be popped.
+ * The stack is modelled by a 64 integer, treating any
+ * stack that can't fit into 64 bits as "overflowed".
+ */
+
 typedef enum kind {
-    With = 1,
-    Loop = 2,
-    Try = 3,
-    Except = 4,
+    Iterator = 1,
+    Except = 2,
+    Object = 3,
 } Kind;
 
-#define BITS_PER_BLOCK 3
+#define BITS_PER_BLOCK 2
+
+#define UNINITIALIZED -2
+#define OVERFLOWED -1
+
+#define MAX_STACK_ENTRIES (63/BITS_PER_BLOCK)
+#define WILL_OVERFLOW (1ULL<<((MAX_STACK_ENTRIES-1)*BITS_PER_BLOCK))
 
 static inline int64_t
-push_block(int64_t stack, Kind kind)
+push_value(int64_t stack, Kind kind)
 {
-    assert(stack < ((int64_t)1)<<(BITS_PER_BLOCK*CO_MAXBLOCKS));
-    return (stack << BITS_PER_BLOCK) | kind;
+    if (((uint64_t)stack) >= WILL_OVERFLOW) {
+        return OVERFLOWED;
+    }
+    else {
+        return (stack << BITS_PER_BLOCK) | kind;
+    }
 }
 
 static inline int64_t
-pop_block(int64_t stack)
+pop_value(int64_t stack)
 {
-    assert(stack > 0);
-    return stack >> BITS_PER_BLOCK;
+    return Py_ARITHMETIC_RIGHT_SHIFT(int64_t, stack, BITS_PER_BLOCK);
 }
 
 static inline Kind
-top_block(int64_t stack)
+top_of_stack(int64_t stack)
 {
     return stack & ((1<<BITS_PER_BLOCK)-1);
 }
 
 static int64_t *
-markblocks(PyCodeObject *code_obj, int len)
+mark_stacks(PyCodeObject *code_obj, int len)
 {
     const _Py_CODEUNIT *code =
         (const _Py_CODEUNIT *)PyBytes_AS_STRING(code_obj->co_code);
-    int64_t *blocks = PyMem_New(int64_t, len+1);
+    int64_t *stacks = PyMem_New(int64_t, len+1);
     int i, j, opcode;
 
-    if (blocks == NULL) {
+    if (stacks == NULL) {
         PyErr_NoMemory();
         return NULL;
     }
-    memset(blocks, -1, (len+1)*sizeof(int64_t));
-    blocks[0] = 0;
+    for (int i = 1; i <= len; i++) {
+        stacks[i] = UNINITIALIZED;
+    }
+    stacks[0] = 0;
     int todo = 1;
     while (todo) {
         todo = 0;
         for (i = 0; i < len; i++) {
-            int64_t block_stack = blocks[i];
-            int64_t except_stack;
-            if (block_stack == -1) {
+            int64_t next_stack = stacks[i];
+            if (next_stack == UNINITIALIZED) {
                 continue;
             }
             opcode = _Py_OPCODE(code[i]);
@@ -150,109 +165,153 @@ markblocks(PyCodeObject *code_obj, int len)
                 case POP_JUMP_IF_FALSE:
                 case POP_JUMP_IF_TRUE:
                 case JUMP_IF_NOT_EXC_MATCH:
-                    j = get_arg(code, i);
+                {
+                    int64_t target_stack;
+                    int j = get_arg(code, i);
                     assert(j < len);
-                    if (blocks[j] == -1 && j < i) {
+                    if (stacks[j] == UNINITIALIZED && j < i) {
                         todo = 1;
                     }
-                    assert(blocks[j] == -1 || blocks[j] == block_stack);
-                    blocks[j] = block_stack;
-                    blocks[i+1] = block_stack;
+                    if (opcode == JUMP_IF_NOT_EXC_MATCH) {
+                        next_stack = pop_value(pop_value(next_stack));
+                        target_stack = next_stack;
+                    }
+                    else if (opcode == JUMP_IF_FALSE_OR_POP ||
+                             opcode == JUMP_IF_TRUE_OR_POP)
+                    {
+                        target_stack = next_stack;
+                        next_stack = pop_value(next_stack);
+                    }
+                    else {
+                        next_stack = pop_value(next_stack);
+                        target_stack = next_stack;
+                    }
+                    assert(stacks[j] == UNINITIALIZED || stacks[j] == target_stack);
+                    stacks[j] = target_stack;
+                    stacks[i+1] = next_stack;
                     break;
+                }
                 case JUMP_ABSOLUTE:
                     j = get_arg(code, i);
                     assert(j < len);
-                    if (blocks[j] == -1 && j < i) {
+                    if (stacks[j] == UNINITIALIZED && j < i) {
                         todo = 1;
                     }
-                    assert(blocks[j] == -1 || blocks[j] == block_stack);
-                    blocks[j] = block_stack;
+                    assert(stacks[j] == UNINITIALIZED || stacks[j] == next_stack);
+                    stacks[j] = next_stack;
                     break;
-                case SETUP_FINALLY:
-                    j = get_arg(code, i) + i + 1;
-                    assert(j < len);
-                    except_stack = push_block(block_stack, Except);
-                    assert(blocks[j] == -1 || blocks[j] == except_stack);
-                    blocks[j] = except_stack;
-                    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);
-                    except_stack = push_block(block_stack, Except);
-                    assert(blocks[j] == -1 || blocks[j] == except_stack);
-                    blocks[j] = except_stack;
-                    block_stack = push_block(block_stack, With);
-                    blocks[i+1] = block_stack;
+                case POP_EXCEPT:
+                    next_stack = pop_value(pop_value(pop_value(next_stack)));
+                    stacks[i+1] = next_stack;
                     break;
+
                 case JUMP_FORWARD:
                     j = get_arg(code, i) + i + 1;
                     assert(j < len);
-                    assert(blocks[j] == -1 || blocks[j] == block_stack);
-                    blocks[j] = block_stack;
+                    assert(stacks[j] == UNINITIALIZED || stacks[j] == next_stack);
+                    stacks[j] = next_stack;
                     break;
                 case GET_ITER:
                 case GET_AITER:
-                    block_stack = push_block(block_stack, Loop);
-                    blocks[i+1] = block_stack;
+                    next_stack = push_value(pop_value(next_stack), Iterator);
+                    stacks[i+1] = next_stack;
                     break;
                 case FOR_ITER:
-                    blocks[i+1] = block_stack;
-                    block_stack = pop_block(block_stack);
+                {
+                    int64_t target_stack = pop_value(next_stack);
+                    stacks[i+1] = push_value(next_stack, Object);
                     j = get_arg(code, i) + i + 1;
                     assert(j < len);
-                    assert(blocks[j] == -1 || blocks[j] == block_stack);
-                    blocks[j] = block_stack;
-                    break;
-                case POP_BLOCK:
-                case POP_EXCEPT:
-                    block_stack = pop_block(block_stack);
-                    blocks[i+1] = block_stack;
+                    assert(stacks[j] == UNINITIALIZED || stacks[j] == target_stack);
+                    stacks[j] = target_stack;
                     break;
+                }
                 case END_ASYNC_FOR:
-                    block_stack = pop_block(pop_block(block_stack));
-                    blocks[i+1] = block_stack;
+                    next_stack = pop_value(pop_value(pop_value(next_stack)));
+                    stacks[i+1] = next_stack;
                     break;
+                case PUSH_EXC_INFO:
+                    next_stack = push_value(next_stack, Except);
+                    next_stack = push_value(next_stack, Except);
+                    next_stack = push_value(next_stack, Except);
+                    stacks[i+1] = next_stack;
                 case RETURN_VALUE:
                 case RAISE_VARARGS:
                 case RERAISE:
+                case POP_EXCEPT_AND_RERAISE:
                     /* End of block */
                     break;
+                case GEN_START:
+                    stacks[i+1] = next_stack;
+                    break;
                 default:
-                    blocks[i+1] = block_stack;
-
+                {
+                    int delta = PyCompile_OpcodeStackEffect(opcode, _Py_OPARG(code[i]));
+                    while (delta < 0) {
+                        next_stack = pop_value(next_stack);
+                        delta++;
+                    }
+                    while (delta > 0) {
+                        next_stack = push_value(next_stack, Object);
+                        delta--;
+                    }
+                    stacks[i+1] = next_stack;
+                }
             }
         }
     }
-    return blocks;
+    return stacks;
+}
+
+static int
+compatible_kind(Kind from, Kind to) {
+    if (to == 0) {
+        return 0;
+    }
+    if (to == Object) {
+        return 1;
+    }
+    return from == to;
 }
 
 static int
-compatible_block_stack(int64_t from_stack, int64_t to_stack)
+compatible_stack(int64_t from_stack, int64_t to_stack)
 {
-    if (to_stack < 0) {
+    if (from_stack < 0 || to_stack < 0) {
         return 0;
     }
     while(from_stack > to_stack) {
-        from_stack = pop_block(from_stack);
+        from_stack = pop_value(from_stack);
     }
-    return from_stack == to_stack;
+    while(from_stack) {
+        Kind from_top = top_of_stack(from_stack);
+        Kind to_top = top_of_stack(to_stack);
+        if (!compatible_kind(from_top, to_top)) {
+            return 0;
+        }
+        from_stack = pop_value(from_stack);
+        to_stack = pop_value(to_stack);
+    }
+    return to_stack == 0;
 }
 
 static const char *
-explain_incompatible_block_stack(int64_t to_stack)
+explain_incompatible_stack(int64_t to_stack)
 {
-    Kind target_kind = top_block(to_stack);
+    assert(to_stack != 0);
+    if (to_stack == OVERFLOWED) {
+        return "stack is too deep to analyze";
+    }
+    if (to_stack == UNINITIALIZED) {
+        return "can't jump into an exception handler, or code may be unreachable";
+    }
+    Kind target_kind = top_of_stack(to_stack);
     switch(target_kind) {
         case Except:
             return "can't jump into an 'except' block as there's no exception";
-        case Try:
-            return "can't jump into the body of a try statement";
-        case With:
-            return "can't jump into the body of a with statement";
-        case Loop:
+        case Object:
+            return "differing stack depth";
+        case Iterator:
             return "can't jump into the body of a for loop";
         default:
             Py_UNREACHABLE();
@@ -299,27 +358,12 @@ first_line_not_before(int *lines, int len, int line)
 static void
 frame_stack_pop(PyFrameObject *f)
 {
-    assert(f->f_stackdepth >= 0);
+    assert(f->f_stackdepth > 0);
     f->f_stackdepth--;
     PyObject *v = f->f_valuestack[f->f_stackdepth];
     Py_DECREF(v);
 }
 
-static void
-frame_block_unwind(PyFrameObject *f)
-{
-    assert(f->f_stackdepth >= 0);
-    assert(f->f_iblock > 0);
-    f->f_iblock--;
-    PyTryBlock *b = &f->f_blockstack[f->f_iblock];
-    intptr_t delta = f->f_stackdepth - b->b_level;
-    while (delta > 0) {
-        frame_stack_pop(f);
-        delta--;
-    }
-}
-
-
 /* 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
@@ -327,13 +371,7 @@ frame_block_unwind(PyFrameObject *f)
  * would still work without any stack errors), but there are some constructs
  * that limit jumping:
  *
- *  o Lines with an 'except' statement on them can't be jumped to, because
- *    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
- *    we cannot be sure which state the interpreter was in or would be in
- *    during execution of the finally block.
- *  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 Any excpetion handlers.
  *  o 'for' and 'async for' loops can't be jumped into because the
  *    iterator needs to be on the stack.
  *  o Jumps cannot be made from within a trace function invoked with a
@@ -428,67 +466,56 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno, void *Py_UNUSED(ignore
         return -1;
     }
 
-    int64_t *blocks = markblocks(f->f_code, len);
-    if (blocks == NULL) {
+    int64_t *stacks = mark_stacks(f->f_code, len);
+    if (stacks == NULL) {
         PyMem_Free(lines);
         return -1;
     }
 
-    int64_t target_block_stack = -1;
-    int64_t best_block_stack = -1;
+    int64_t best_stack = OVERFLOWED;
     int best_addr = -1;
-    int64_t start_block_stack = blocks[f->f_lasti];
+    int64_t start_stack = stacks[f->f_lasti];
+    int err = -1;
     const char *msg = "cannot find bytecode for specified line";
     for (int i = 0; i < len; i++) {
         if (lines[i] == new_lineno) {
-            target_block_stack = blocks[i];
-            if (compatible_block_stack(start_block_stack, target_block_stack)) {
-                msg = NULL;
-                if (target_block_stack > best_block_stack) {
-                    best_block_stack = target_block_stack;
+            int64_t target_stack = stacks[i];
+            if (compatible_stack(start_stack, target_stack)) {
+                err = 0;
+                if (target_stack > best_stack) {
+                    best_stack = target_stack;
                     best_addr = i;
                 }
             }
-            else if (msg) {
-                if (target_block_stack >= 0) {
-                    msg = explain_incompatible_block_stack(target_block_stack);
+            else if (err < 0) {
+                if (start_stack == OVERFLOWED) {
+                    msg = "stack to deep to analyze";
+                }
+                else if (start_stack == UNINITIALIZED) {
+                    msg = "can't jump from within an exception handler";
                 }
                 else {
-                    msg = "code may be unreachable.";
+                    msg = explain_incompatible_stack(target_stack);
+                    err = 1;
                 }
             }
         }
     }
-    PyMem_Free(blocks);
+    PyMem_Free(stacks);
     PyMem_Free(lines);
-    if (msg != NULL) {
+    if (err) {
         PyErr_SetString(PyExc_ValueError, msg);
         return -1;
     }
-
     /* Unwind block stack. */
-    while (start_block_stack > best_block_stack) {
-        Kind kind = top_block(start_block_stack);
-        switch(kind) {
-        case Loop:
-            frame_stack_pop(f);
-            break;
-        case Try:
-            frame_block_unwind(f);
-            break;
-        case With:
-            frame_block_unwind(f);
-            // Pop the exit function
-            frame_stack_pop(f);
-            break;
-        case Except:
-            PyErr_SetString(PyExc_ValueError,
-                "can't jump out of an 'except' block");
-            return -1;
-        }
-        start_block_stack = pop_block(start_block_stack);
+    if (f->f_state == FRAME_SUSPENDED) {
+        /* Account for value popped by yield */
+        start_stack = pop_value(start_stack);
+    }
+    while (start_stack > best_stack) {
+        frame_stack_pop(f);
+        start_stack = pop_value(start_stack);
     }
-
     /* Finally set the new f_lasti and return OK. */
     f->f_lineno = 0;
     f->f_lasti = best_addr;
@@ -852,7 +879,6 @@ _PyFrame_New_NoTrack(PyThreadState *tstate, PyFrameConstructor *con, PyObject *l
     f->f_gen = NULL;
     f->f_lasti = -1;
     f->f_lineno = 0;
-    f->f_iblock = 0;
     f->f_state = FRAME_CREATED;
     // f_blockstack and f_localsplus initialized by frame_alloc()
     return f;
@@ -884,33 +910,6 @@ PyFrame_New(PyThreadState *tstate, PyCodeObject *code,
     return f;
 }
 
-
-/* Block management */
-
-void
-PyFrame_BlockSetup(PyFrameObject *f, int type, int handler, int level)
-{
-    PyTryBlock *b;
-    if (f->f_iblock >= CO_MAXBLOCKS) {
-        Py_FatalError("block stack overflow");
-    }
-    b = &f->f_blockstack[f->f_iblock++];
-    b->b_type = type;
-    b->b_level = level;
-    b->b_handler = handler;
-}
-
-PyTryBlock *
-PyFrame_BlockPop(PyFrameObject *f)
-{
-    PyTryBlock *b;
-    if (f->f_iblock <= 0) {
-        Py_FatalError("block stack underflow");
-    }
-    b = &f->f_blockstack[--f->f_iblock];
-    return b;
-}
-
 /* Convert between "fast" version of locals and dictionary version.
 
    map and values are input arguments.  map is a tuple of strings.