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>

Doc/library/dis.rst

@@ -335,6 +335,14 @@ The Python compiler currently generates the following bytecode instructions.
    three.
 
 
+.. opcode:: ROT_FOUR
+
+   Lifts second, third and forth stack items one position up, moves top down
+   to position four.
+
+   .. versionadded:: 3.8
+
+
 .. opcode:: DUP_TOP
 
    Duplicates the reference on top of the stack.
@@ -605,17 +613,6 @@ the original TOS1.
    is terminated with :opcode:`POP_TOP`.
 
 
-.. opcode:: BREAK_LOOP
-
-   Terminates a loop due to a :keyword:`break` statement.
-
-
-.. opcode:: CONTINUE_LOOP (target)
-
-   Continues a loop due to a :keyword:`continue` statement.  *target* is the
-   address to jump to (which should be a :opcode:`FOR_ITER` instruction).
-
-
 .. opcode:: SET_ADD (i)
 
    Calls ``set.add(TOS1[-i], TOS)``.  Used to implement set comprehensions.
@@ -676,7 +673,7 @@ iterations of the loop.
 .. opcode:: POP_BLOCK
 
    Removes one block from the block stack.  Per frame, there is a stack of
-   blocks, denoting nested loops, try statements, and such.
+   blocks, denoting :keyword:`try` statements, and such.
 
 
 .. opcode:: POP_EXCEPT
@@ -687,11 +684,50 @@ iterations of the loop.
    popped values are used to restore the exception state.
 
 
+.. opcode:: POP_FINALLY (preserve_tos)
+
+   Cleans up the value stack and the block stack.  If *preserve_tos* is not
+   ``0`` TOS first is popped from the stack and pushed on the stack after
+   perfoming other stack operations:
+
+   * If TOS is ``NULL`` or an integer (pushed by :opcode:`BEGIN_FINALLY`
+     or :opcode:`CALL_FINALLY`) it is popped from the stack.
+   * If TOS is an exception type (pushed when an exception has been raised)
+     6 values are popped from the stack, the last three popped values are
+     used to restore the exception state.  An exception handler block is
+     removed from the block stack.
+
+   It is similar to :opcode:`END_FINALLY`, but doesn't change the bytecode
+   counter nor raise an exception.  Used for implementing :keyword:`break`
+   and :keyword:`return` in the :keyword:`finally` block.
+
+   .. versionadded:: 3.8
+
+
+.. opcode:: BEGIN_FINALLY
+
+   Pushes ``NULL`` onto the stack for using it in :opcode:`END_FINALLY`,
+   :opcode:`POP_FINALLY`, :opcode:`WITH_CLEANUP_START` and
+   :opcode:`WITH_CLEANUP_FINISH`.  Starts the :keyword:`finally` block.
+
+   .. versionadded:: 3.8
+
+
 .. opcode:: END_FINALLY
 
    Terminates a :keyword:`finally` clause.  The interpreter recalls whether the
-   exception has to be re-raised, or whether the function returns, and continues
-   with the outer-next block.
+   exception has to be re-raised or execution has to be continued depending on
+   the value of TOS.
+
+   * If TOS is ``NULL`` (pushed by :opcode:`BEGIN_FINALLY`) continue from
+     the next instruction. TOS is popped.
+   * If TOS is an integer (pushed by :opcode:`CALL_FINALLY`), sets the
+     bytecode counter to TOS.  TOS is popped.
+   * If TOS is an exception type (pushed when an exception has been raised)
+     6 values are popped from the stack, the first three popped values are
+     used to re-raise the exception and the last three popped values are used
+     to restore the exception state.  An exception handler block is removed
+     from the block stack.
 
 
 .. opcode:: LOAD_BUILD_CLASS
@@ -704,9 +740,9 @@ iterations of the loop.
 
    This opcode performs several operations before a with block starts.  First,
    it loads :meth:`~object.__exit__` from the context manager and pushes it onto
-   the stack for later use by :opcode:`WITH_CLEANUP`.  Then,
+   the stack for later use by :opcode:`WITH_CLEANUP_START`.  Then,
    :meth:`~object.__enter__` is called, and a finally block pointing to *delta*
-   is pushed.  Finally, the result of calling the enter method is pushed onto
+   is pushed.  Finally, the result of calling the ``__enter__()`` method is pushed onto
    the stack.  The next opcode will either ignore it (:opcode:`POP_TOP`), or
    store it in (a) variable(s) (:opcode:`STORE_FAST`, :opcode:`STORE_NAME`, or
    :opcode:`UNPACK_SEQUENCE`).
@@ -716,30 +752,31 @@ iterations of the loop.
 
 .. opcode:: WITH_CLEANUP_START
 
-   Cleans up the stack when a :keyword:`with` statement block exits.  TOS is the
-   context manager's :meth:`__exit__` bound method. Below TOS are 1--3 values
-   indicating how/why the finally clause was entered:
+   Starts cleaning up the stack when a :keyword:`with` statement block exits.
 
-   * SECOND = ``None``
-   * (SECOND, THIRD) = (``WHY_{RETURN,CONTINUE}``), retval
-   * SECOND = ``WHY_*``; no retval below it
-   * (SECOND, THIRD, FOURTH) = exc_info()
+   At the top of the stack are either ``NULL`` (pushed by
+   :opcode:`BEGIN_FINALLY`) or 6 values pushed if an exception has been
+   raised in the with block.  Below is the context manager's
+   :meth:`~object.__exit__` or :meth:`~object.__aexit__` bound method.
 
-   In the last case, ``TOS(SECOND, THIRD, FOURTH)`` is called, otherwise
-   ``TOS(None, None, None)``.  Pushes SECOND and result of the call
-   to the stack.
+   If TOS is ``NULL``, calls ``SECOND(None, None, None)``,
+   removes the function from the stack, leaving TOS, and pushes ``None``
+   to the stack.  Otherwise calls ``SEVENTH(TOP, SECOND, THIRD)``,
+   shifts the bottom 3 values of the stack down, replaces the empty spot
+   with ``NULL`` and pushes TOS.  Finally pushes the result of the call.
 
 
 .. opcode:: WITH_CLEANUP_FINISH
 
-   Pops exception type and result of 'exit' function call from the stack.
+   Finishes cleaning up the stack when a :keyword:`with` statement block exits.
 
-   If the stack represents an exception, *and* the function call returns a
-   'true' value, this information is "zapped" and replaced with a single
-   ``WHY_SILENCED`` to prevent :opcode:`END_FINALLY` from re-raising the
-   exception.  (But non-local gotos will still be resumed.)
+   TOS is result of ``__exit__()`` or ``__aexit__()`` function call pushed
+   by :opcode:`WITH_CLEANUP_START`.  SECOND is ``None`` or an exception type
+   (pushed when an exception has been raised).
 
-   .. XXX explain the WHY stuff!
+   Pops two values from the stack.  If SECOND is not None and TOS is true
+   unwinds the EXCEPT_HANDLER block which was created when the exception
+   was caught and pushes ``NULL`` to the stack.
 
 
 All of the following opcodes use their arguments.
@@ -987,22 +1024,19 @@ All of the following opcodes use their arguments.
    Loads the global named ``co_names[namei]`` onto the stack.
 
 
-.. opcode:: SETUP_LOOP (delta)
-
-   Pushes a block for a loop onto the block stack.  The block spans from the
-   current instruction with a size of *delta* bytes.
-
+.. opcode:: SETUP_FINALLY (delta)
 
-.. opcode:: SETUP_EXCEPT (delta)
+   Pushes a try block from a try-finally or try-except clause onto the block
+   stack.  *delta* points to the finally block or the first except block.
 
-   Pushes a try block from a try-except clause onto the block stack. *delta*
-   points to the first except block.
 
+.. opcode:: CALL_FINALLY (delta)
 
-.. opcode:: SETUP_FINALLY (delta)
+   Pushes the address of the next instruction onto the stack and increments
+   bytecode counter by *delta*.  Used for calling the finally block as a
+   "subroutine".
 
-   Pushes a try block from a try-except clause onto the block stack. *delta*
-   points to the finally block.
+   .. versionadded:: 3.8
 
 
 .. opcode:: LOAD_FAST (var_num)

Doc/whatsnew/3.8.rst

@@ -137,3 +137,21 @@ Changes in the Python API
   :func:`dbm.dumb.open` with flags ``'r'`` and ``'w'`` no longer creates
   a database if it does not exist.
   (Contributed by Serhiy Storchaka in :issue:`32749`.)
+
+
+CPython bytecode changes
+------------------------
+
+* The interpreter loop  has been simplified by moving the logic of unrolling
+  the stack of blocks into the compiler.  The compiler emits now explicit
+  instructions for adjusting the stack of values and calling the cleaning
+  up code for :keyword:`break`, :keyword:`continue` and :keyword:`return`.
+
+  Removed opcodes :opcode:`BREAK_LOOP`, :opcode:`CONTINUE_LOOP`,
+  :opcode:`SETUP_LOOP` and :opcode:`SETUP_EXCEPT`.  Added new opcodes
+  :opcode:`ROT_FOUR`, :opcode:`BEGIN_FINALLY`, :opcode:`CALL_FINALLY` and
+  :opcode:`POP_FINALLY`.  Changed the behavior of :opcode:`END_FINALLY`
+  and :opcode:`WITH_CLEANUP_START`.
+
+  (Contributed by Mark Shannon, Antoine Pitrou and Serhiy Storchaka in
+  :issue:`17611`.)

Include/opcode.h

@@ -12,6 +12,7 @@ extern "C" {
 #define ROT_THREE                 3
 #define DUP_TOP                   4
 #define DUP_TOP_TWO               5
+#define ROT_FOUR                  6
 #define NOP                       9
 #define UNARY_POSITIVE           10
 #define UNARY_NEGATIVE           11
@@ -32,6 +33,7 @@ extern "C" {
 #define GET_AITER                50
 #define GET_ANEXT                51
 #define BEFORE_ASYNC_WITH        52
+#define BEGIN_FINALLY            53
 #define INPLACE_ADD              55
 #define INPLACE_SUBTRACT         56
 #define INPLACE_MULTIPLY         57
@@ -55,7 +57,6 @@ extern "C" {
 #define INPLACE_AND              77
 #define INPLACE_XOR              78
 #define INPLACE_OR               79
-#define BREAK_LOOP               80
 #define WITH_CLEANUP_START       81
 #define WITH_CLEANUP_FINISH      82
 #define RETURN_VALUE             83
@@ -92,9 +93,6 @@ extern "C" {
 #define POP_JUMP_IF_FALSE       114
 #define POP_JUMP_IF_TRUE        115
 #define LOAD_GLOBAL             116
-#define CONTINUE_LOOP           119
-#define SETUP_LOOP              120
-#define SETUP_EXCEPT            121
 #define SETUP_FINALLY           122
 #define LOAD_FAST               124
 #define STORE_FAST              125
@@ -127,6 +125,8 @@ extern "C" {
 #define BUILD_TUPLE_UNPACK_WITH_CALL 158
 #define LOAD_METHOD             160
 #define CALL_METHOD             161
+#define CALL_FINALLY            162
+#define POP_FINALLY             163
 
 /* EXCEPT_HANDLER is a special, implicit block type which is created when
    entering an except handler. It is not an opcode but we define it here

Lib/importlib/_bootstrap_external.py

@@ -246,6 +246,7 @@ _code_type = type(_write_atomic.__code__)
 #     Python 3.7a2  3391 (update GET_AITER #31709)
 #     Python 3.7a4  3392 (PEP 552: Deterministic pycs #31650)
 #     Python 3.7b1  3393 (remove STORE_ANNOTATION opcode #32550)
+#     Python 3.8a1  3400 (move frame block handling to compiler #17611)
 #
 # MAGIC must change whenever the bytecode emitted by the compiler may no
 # longer be understood by older implementations of the eval loop (usually
@@ -254,7 +255,7 @@ _code_type = type(_write_atomic.__code__)
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3393).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3400).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'

Lib/opcode.py

@@ -60,6 +60,7 @@ def_op('ROT_TWO', 2)
 def_op('ROT_THREE', 3)
 def_op('DUP_TOP', 4)
 def_op('DUP_TOP_TWO', 5)
+def_op('ROT_FOUR', 6)
 
 def_op('NOP', 9)
 def_op('UNARY_POSITIVE', 10)
@@ -86,6 +87,7 @@ def_op('INPLACE_TRUE_DIVIDE', 29)
 def_op('GET_AITER', 50)
 def_op('GET_ANEXT', 51)
 def_op('BEFORE_ASYNC_WITH', 52)
+def_op('BEGIN_FINALLY', 53)
 
 def_op('INPLACE_ADD', 55)
 def_op('INPLACE_SUBTRACT', 56)
@@ -113,10 +115,8 @@ def_op('INPLACE_RSHIFT', 76)
 def_op('INPLACE_AND', 77)
 def_op('INPLACE_XOR', 78)
 def_op('INPLACE_OR', 79)
-def_op('BREAK_LOOP', 80)
 def_op('WITH_CLEANUP_START', 81)
 def_op('WITH_CLEANUP_FINISH', 82)
-
 def_op('RETURN_VALUE', 83)
 def_op('IMPORT_STAR', 84)
 def_op('SETUP_ANNOTATIONS', 85)
@@ -158,10 +158,7 @@ jabs_op('POP_JUMP_IF_TRUE', 115)     # ""
 
 name_op('LOAD_GLOBAL', 116)     # Index in name list
 
-jabs_op('CONTINUE_LOOP', 119)   # Target address
-jrel_op('SETUP_LOOP', 120)      # Distance to target address
-jrel_op('SETUP_EXCEPT', 121)    # ""
-jrel_op('SETUP_FINALLY', 122)   # ""
+jrel_op('SETUP_FINALLY', 122)   # Distance to target address
 
 def_op('LOAD_FAST', 124)        # Local variable number
 haslocal.append(124)
@@ -213,5 +210,7 @@ def_op('BUILD_TUPLE_UNPACK_WITH_CALL', 158)
 
 name_op('LOAD_METHOD', 160)
 def_op('CALL_METHOD', 161)
+jrel_op('CALL_FINALLY', 162)
+def_op('POP_FINALLY', 163)
 
 del def_op, name_op, jrel_op, jabs_op

Lib/test/test_importlib/test_util.py

@@ -789,7 +789,7 @@ class MagicNumberTests(unittest.TestCase):
         in advance. Such exceptional releases will then require an
         adjustment to this test case.
         """
-        EXPECTED_MAGIC_NUMBER = 3379
+        EXPECTED_MAGIC_NUMBER = 3400
         actual = int.from_bytes(importlib.util.MAGIC_NUMBER[:2], 'little')
 
         msg = (

Lib/test/test_peepholer.py

@@ -268,7 +268,7 @@ class TestTranforms(BytecodeTestCase):
         self.assertNotInBytecode(f, 'JUMP_ABSOLUTE')
         returns = [instr for instr in dis.get_instructions(f)
                           if instr.opname == 'RETURN_VALUE']
-        self.assertEqual(len(returns), 6)
+        self.assertLessEqual(len(returns), 6)
 
     def test_elim_jump_after_return2(self):
         # Eliminate dead code: jumps immediately after returns can't be reached
@@ -282,7 +282,7 @@ class TestTranforms(BytecodeTestCase):
         self.assertEqual(len(returns), 1)
         returns = [instr for instr in dis.get_instructions(f)
                           if instr.opname == 'RETURN_VALUE']
-        self.assertEqual(len(returns), 2)
+        self.assertLessEqual(len(returns), 2)
 
     def test_make_function_doesnt_bail(self):
         def f():

Lib/test/test_sys_settrace.py

@@ -187,7 +187,6 @@ tightloop_example.events = [(0, 'call'),
                             (1, 'line'),
                             (2, 'line'),
                             (3, 'line'),
-                            (4, 'line'),
                             (5, 'line'),
                             (5, 'line'),
                             (5, 'line'),
@@ -715,6 +714,21 @@ class JumpTestCase(unittest.TestCase):
             output.append(11)
         output.append(12)
 
+    @jump_test(5, 11, [2, 4, 12])
+    def test_jump_over_return_try_finally_in_finally_block(output):
+        try:
+            output.append(2)
+        finally:
+            output.append(4)
+            output.append(5)
+            return
+            try:
+                output.append(8)
+            finally:
+                output.append(10)
+            pass
+        output.append(12)
+
     @jump_test(3, 4, [1, 4])
     def test_jump_infinite_while_loop(output):
         output.append(1)
@@ -722,6 +736,22 @@ class JumpTestCase(unittest.TestCase):
             output.append(3)
         output.append(4)
 
+    @jump_test(2, 4, [4, 4])
+    def test_jump_forwards_into_while_block(output):
+        i = 1
+        output.append(2)
+        while i <= 2:
+            output.append(4)
+            i += 1
+
+    @jump_test(5, 3, [3, 3, 3, 5])
+    def test_jump_backwards_into_while_block(output):
+        i = 1
+        while i <= 2:
+            output.append(3)
+            i += 1
+        output.append(5)
+
     @jump_test(2, 3, [1, 3])
     def test_jump_forwards_out_of_with_block(output):
         with tracecontext(output, 1):
@@ -916,22 +946,6 @@ class JumpTestCase(unittest.TestCase):
             output.append(2)
         output.append(3)
 
-    @jump_test(2, 4, [], (ValueError, 'into'))
-    def test_no_jump_forwards_into_while_block(output):
-        i = 1
-        output.append(2)
-        while i <= 2:
-            output.append(4)
-            i += 1
-
-    @jump_test(5, 3, [3, 3], (ValueError, 'into'))
-    def test_no_jump_backwards_into_while_block(output):
-        i = 1
-        while i <= 2:
-            output.append(3)
-            i += 1
-        output.append(5)
-
     @jump_test(1, 3, [], (ValueError, 'into'))
     def test_no_jump_forwards_into_with_block(output):
         output.append(1)
@@ -1024,6 +1038,16 @@ class JumpTestCase(unittest.TestCase):
         with tracecontext(output, 4):
             output.append(5)
 
+    @jump_test(5, 7, [2, 4], (ValueError, 'finally'))
+    def test_no_jump_over_return_out_of_finally_block(output):
+        try:
+            output.append(2)
+        finally:
+            output.append(4)
+            output.append(5)
+            return
+        output.append(7)
+
     @jump_test(7, 4, [1, 6], (ValueError, 'into'))
     def test_no_jump_into_for_block_before_else(output):
         output.append(1)

Misc/NEWS.d/next/Core and Builtins/2017-12-24-19-48-59.bpo-17611.P85kWL.rst

+Simplified the interpreter loop by moving the logic of unrolling the stack
+of blocks into the compiler. The compiler emits now explicit instructions
+for adjusting the stack of values and calling the cleaning up code for
+:keyword:`break`, :keyword:`continue` and :keyword:`return`.
+
+Removed opcodes :opcode:`BREAK_LOOP`, :opcode:`CONTINUE_LOOP`,
+:opcode:`SETUP_LOOP` and :opcode:`SETUP_EXCEPT`. Added new opcodes
+:opcode:`ROT_FOUR`, :opcode:`BEGIN_FINALLY` and :opcode:`CALL_FINALLY` and
+:opcode:`POP_FINALLY`. Changed the behavior of :opcode:`END_FINALLY` and
+:opcode:`WITH_CLEANUP_START`.

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;
 }
 

PC/launcher.c

@@ -1116,6 +1116,7 @@ static PYC_MAGIC magic_values[] = {
     { 3320, 3351, L"3.5" },
     { 3360, 3379, L"3.6" },
     { 3390, 3399, L"3.7" },
+    { 3400, 3409, L"3.8" },
     { 0 }
 };
 

Python/opcode_targets.h

@@ -5,7 +5,7 @@ static void *opcode_targets[256] = {
     &&TARGET_ROT_THREE,
     &&TARGET_DUP_TOP,
     &&TARGET_DUP_TOP_TWO,
-    &&_unknown_opcode,
+    &&TARGET_ROT_FOUR,
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&TARGET_NOP,
@@ -52,7 +52,7 @@ static void *opcode_targets[256] = {
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
     &&TARGET_BEFORE_ASYNC_WITH,
-    &&_unknown_opcode,
+    &&TARGET_BEGIN_FINALLY,
     &&_unknown_opcode,
     &&TARGET_INPLACE_ADD,
     &&TARGET_INPLACE_SUBTRACT,
@@ -79,7 +79,7 @@ static void *opcode_targets[256] = {
     &&TARGET_INPLACE_AND,
     &&TARGET_INPLACE_XOR,
     &&TARGET_INPLACE_OR,
-    &&TARGET_BREAK_LOOP,
+    &&_unknown_opcode,
     &&TARGET_WITH_CLEANUP_START,
     &&TARGET_WITH_CLEANUP_FINISH,
     &&TARGET_RETURN_VALUE,
@@ -118,9 +118,9 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_GLOBAL,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&TARGET_CONTINUE_LOOP,
-    &&TARGET_SETUP_LOOP,
-    &&TARGET_SETUP_EXCEPT,
+    &&_unknown_opcode,
+    &&_unknown_opcode,
+    &&_unknown_opcode,
     &&TARGET_SETUP_FINALLY,
     &&_unknown_opcode,
     &&TARGET_LOAD_FAST,
@@ -161,8 +161,8 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&TARGET_LOAD_METHOD,
     &&TARGET_CALL_METHOD,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
+    &&TARGET_CALL_FINALLY,
+    &&TARGET_POP_FINALLY,
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,

Python/peephole.c

@@ -13,7 +13,7 @@
 #define UNCONDITIONAL_JUMP(op)  (op==JUMP_ABSOLUTE || op==JUMP_FORWARD)
 #define CONDITIONAL_JUMP(op) (op==POP_JUMP_IF_FALSE || op==POP_JUMP_IF_TRUE \
     || op==JUMP_IF_FALSE_OR_POP || op==JUMP_IF_TRUE_OR_POP)
-#define ABSOLUTE_JUMP(op) (op==JUMP_ABSOLUTE || op==CONTINUE_LOOP \
+#define ABSOLUTE_JUMP(op) (op==JUMP_ABSOLUTE \
     || op==POP_JUMP_IF_FALSE || op==POP_JUMP_IF_TRUE \
     || op==JUMP_IF_FALSE_OR_POP || op==JUMP_IF_TRUE_OR_POP)
 #define JUMPS_ON_TRUE(op) (op==POP_JUMP_IF_TRUE || op==JUMP_IF_TRUE_OR_POP)
@@ -185,12 +185,10 @@ markblocks(_Py_CODEUNIT *code, Py_ssize_t len)
             case POP_JUMP_IF_FALSE:
             case POP_JUMP_IF_TRUE:
             case JUMP_ABSOLUTE:
-            case CONTINUE_LOOP:
-            case SETUP_LOOP:
-            case SETUP_EXCEPT:
             case SETUP_FINALLY:
             case SETUP_WITH:
             case SETUP_ASYNC_WITH:
+            case CALL_FINALLY:
                 j = GETJUMPTGT(code, i);
                 assert(j < len);
                 blocks[j] = 1;
@@ -373,15 +371,8 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
                 /* Replace jumps to unconditional jumps */
             case POP_JUMP_IF_FALSE:
             case POP_JUMP_IF_TRUE:
-            case FOR_ITER:
             case JUMP_FORWARD:
             case JUMP_ABSOLUTE:
-            case CONTINUE_LOOP:
-            case SETUP_LOOP:
-            case SETUP_EXCEPT:
-            case SETUP_FINALLY:
-            case SETUP_WITH:
-            case SETUP_ASYNC_WITH:
                 h = GETJUMPTGT(codestr, i);
                 tgt = find_op(codestr, h);
                 /* Replace JUMP_* to a RETURN into just a RETURN */
@@ -407,7 +398,21 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
                 /* Remove unreachable ops after RETURN */
             case RETURN_VALUE:
                 h = i + 1;
-                while (h < codelen && ISBASICBLOCK(blocks, i, h)) {
+                /* END_FINALLY should be kept since it denotes the end of
+                   the 'finally' block in frame_setlineno() in frameobject.c.
+                   SETUP_FINALLY should be kept for balancing.
+                 */
+                while (h < codelen && ISBASICBLOCK(blocks, i, h) &&
+                       _Py_OPCODE(codestr[h]) != END_FINALLY)
+                {
+                    if (_Py_OPCODE(codestr[h]) == SETUP_FINALLY) {
+                        while (h > i + 1 &&
+                               _Py_OPCODE(codestr[h - 1]) == EXTENDED_ARG)
+                        {
+                            h--;
+                        }
+                        break;
+                    }
                     h++;
                 }
                 if (h > i + 1) {
@@ -452,7 +457,6 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
             case NOP:continue;
 
             case JUMP_ABSOLUTE:
-            case CONTINUE_LOOP:
             case POP_JUMP_IF_FALSE:
             case POP_JUMP_IF_TRUE:
             case JUMP_IF_FALSE_OR_POP:
@@ -462,11 +466,10 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
 
             case FOR_ITER:
             case JUMP_FORWARD:
-            case SETUP_LOOP:
-            case SETUP_EXCEPT:
             case SETUP_FINALLY:
             case SETUP_WITH:
             case SETUP_ASYNC_WITH:
+            case CALL_FINALLY:
                 j = blocks[j / sizeof(_Py_CODEUNIT) + i + 1] - blocks[i] - 1;
                 j *= sizeof(_Py_CODEUNIT);
                 break;