bpo-42246: Don't eliminate jumps to jumps, if it will break PEP 626. (GH-23896)

1 files changed, 84 insertions, 51 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 6698b55..4ba9140 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1432,28 +1432,32 @@ compiler_addop_i(struct compiler *c, int opcode, Py_ssize_t oparg)
     return 1;
 }
 
-static int
-compiler_addop_j(struct compiler *c, int opcode, basicblock *b)
+static int add_jump_to_block(basicblock *b, int opcode, int lineno, basicblock *target)
 {
-    struct instr *i;
-    int off;
-
-    if (c->c_do_not_emit_bytecode) {
-        return 1;
-    }
-
     assert(HAS_ARG(opcode));
     assert(b != NULL);
-    off = compiler_next_instr(c->u->u_curblock);
-    if (off < 0)
+    assert(target != NULL);
+
+    int off = compiler_next_instr(b);
+    struct instr *i = &b->b_instr[off];
+    if (off < 0) {
         return 0;
-    i = &c->u->u_curblock->b_instr[off];
+    }
     i->i_opcode = opcode;
-    i->i_target = b;
-    i->i_lineno = c->u->u_lineno;
+    i->i_target = target;
+    i->i_lineno = lineno;
     return 1;
 }
 
+static int
+compiler_addop_j(struct compiler *c, int opcode, basicblock *b)
+{
+    if (c->c_do_not_emit_bytecode) {
+        return 1;
+    }
+    return add_jump_to_block(c->u->u_curblock, opcode, c->u->u_lineno, b);
+}
+
 /* NEXT_BLOCK() creates an implicit jump from the current block
    to the new block.
 
@@ -6067,6 +6071,27 @@ fold_tuple_on_constants(struct instr *inst,
     return 0;
 }
 
+
+static int
+eliminate_jump_to_jump(basicblock *bb, int opcode) {
+    assert (bb->b_iused > 0);
+    struct instr *inst = &bb->b_instr[bb->b_iused-1];
+    assert (is_jump(inst));
+    assert (inst->i_target->b_iused > 0);
+    struct instr *target = &inst->i_target->b_instr[0];
+    if (inst->i_target == target->i_target) {
+        /* Nothing to do */
+        return 0;
+    }
+    int lineno = target->i_lineno;
+    if (add_jump_to_block(bb, opcode, lineno, target->i_target) == 0) {
+        return -1;
+    }
+    assert (bb->b_iused >= 2);
+    bb->b_instr[bb->b_iused-2].i_opcode = NOP;
+    return 0;
+}
+
 /* Maximum size of basic block that should be copied in optimizer */
 #define MAX_COPY_SIZE 4
 
@@ -6183,22 +6208,27 @@ optimize_basic_block(basicblock *bb, PyObject *consts)
             case JUMP_IF_FALSE_OR_POP:
                 switch(target->i_opcode) {
                     case POP_JUMP_IF_FALSE:
-                        *inst = *target;
-                        --i;
+                        if (inst->i_lineno == target->i_lineno) {
+                            *inst = *target;
+                            i--;
+                        }
                         break;
                     case JUMP_ABSOLUTE:
                     case JUMP_FORWARD:
                     case JUMP_IF_FALSE_OR_POP:
-                        if (inst->i_target != target->i_target) {
+                        if (inst->i_lineno == target->i_lineno &&
+                            inst->i_target != target->i_target) {
                             inst->i_target = target->i_target;
-                            --i;
+                            i--;
                         }
                         break;
                     case JUMP_IF_TRUE_OR_POP:
                         assert (inst->i_target->b_iused == 1);
-                        inst->i_opcode = POP_JUMP_IF_FALSE;
-                        inst->i_target = inst->i_target->b_next;
-                        --i;
+                        if (inst->i_lineno == target->i_lineno) {
+                            inst->i_opcode = POP_JUMP_IF_FALSE;
+                            inst->i_target = inst->i_target->b_next;
+                            --i;
+                        }
                         break;
                 }
                 break;
@@ -6206,22 +6236,27 @@ optimize_basic_block(basicblock *bb, PyObject *consts)
             case JUMP_IF_TRUE_OR_POP:
                 switch(target->i_opcode) {
                     case POP_JUMP_IF_TRUE:
-                        *inst = *target;
-                        --i;
+                        if (inst->i_lineno == target->i_lineno) {
+                            *inst = *target;
+                            i--;
+                        }
                         break;
                     case JUMP_ABSOLUTE:
                     case JUMP_FORWARD:
                     case JUMP_IF_TRUE_OR_POP:
-                        if (inst->i_target != target->i_target) {
+                        if (inst->i_lineno == target->i_lineno &&
+                            inst->i_target != target->i_target) {
                             inst->i_target = target->i_target;
-                            --i;
+                            i--;
                         }
                         break;
                     case JUMP_IF_FALSE_OR_POP:
                         assert (inst->i_target->b_iused == 1);
-                        inst->i_opcode = POP_JUMP_IF_TRUE;
-                        inst->i_target = inst->i_target->b_next;
-                        --i;
+                        if (inst->i_lineno == target->i_lineno) {
+                            inst->i_opcode = POP_JUMP_IF_TRUE;
+                            inst->i_target = inst->i_target->b_next;
+                            --i;
+                        }
                         break;
                 }
                 break;
@@ -6230,9 +6265,9 @@ optimize_basic_block(basicblock *bb, PyObject *consts)
                 switch(target->i_opcode) {
                     case JUMP_ABSOLUTE:
                     case JUMP_FORWARD:
-                        if (inst->i_target != target->i_target) {
+                        if (inst->i_lineno == target->i_lineno) {
                             inst->i_target = target->i_target;
-                            --i;
+                            i--;
                         }
                         break;
                 }
@@ -6242,9 +6277,9 @@ optimize_basic_block(basicblock *bb, PyObject *consts)
                 switch(target->i_opcode) {
                     case JUMP_ABSOLUTE:
                     case JUMP_FORWARD:
-                        if (inst->i_target != target->i_target) {
+                        if (inst->i_lineno == target->i_lineno) {
                             inst->i_target = target->i_target;
-                            --i;
+                            i--;
                         }
                         break;
                 }
@@ -6255,32 +6290,30 @@ optimize_basic_block(basicblock *bb, PyObject *consts)
                 assert (i == bb->b_iused-1);
                 switch(target->i_opcode) {
                     case JUMP_FORWARD:
-                        if (inst->i_target != target->i_target) {
-                            inst->i_target = target->i_target;
-//                             --i;
+                        if (eliminate_jump_to_jump(bb, inst->i_opcode)) {
+                            goto error;
                         }
                         break;
+
                     case JUMP_ABSOLUTE:
-                        if (inst->i_target != target->i_target) {
-                            inst->i_target = target->i_target;
-                            inst->i_opcode = target->i_opcode;
-                            --i;
+                        if (eliminate_jump_to_jump(bb, JUMP_ABSOLUTE)) {
+                            goto error;
                         }
                         break;
-                }
-                if (inst->i_target->b_exit && inst->i_target->b_iused <= MAX_COPY_SIZE) {
-                    basicblock *to_copy = inst->i_target;
-                    inst->i_opcode = NOP;
-                    for (i = 0; i < to_copy->b_iused; i++) {
-                        int index = compiler_next_instr(bb);
-                        if (index < 0) {
-                            return -1;
+                    default:
+                        if (inst->i_target->b_exit && inst->i_target->b_iused <= MAX_COPY_SIZE) {
+                            basicblock *to_copy = inst->i_target;
+                            inst->i_opcode = NOP;
+                            for (i = 0; i < to_copy->b_iused; i++) {
+                                int index = compiler_next_instr(bb);
+                                if (index < 0) {
+                                    return -1;
+                                }
+                                bb->b_instr[index] = to_copy->b_instr[i];
+                            }
+                            bb->b_exit = 1;
                         }
-                        bb->b_instr[index] = to_copy->b_instr[i];
-                    }
-                    bb->b_exit = 1;
                 }
-                break;
         }
     }
     return 0;