From 894cafd9a53f2423adc6b1fb111006883aeac03c Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Thu, 1 Sep 2022 11:03:52 +0100
Subject: [PATCH] gh-93678: apply remove_redundant_jumps in optimize_cfg
 (GH-96274)

---
 Lib/test/test_peepholer.py |  6 +++---
 Python/compile.c           | 44 ++++++++++++++++++++++++++++----------
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
index 7ece468363be58..7cd9b060f5e90b 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -892,7 +892,8 @@ def test_conditional_jump_forward_non_const_condition(self):
         self.cfg_optimization_test(insts, expected, consts=list(range(5)))
 
     def test_conditional_jump_forward_const_condition(self):
-        # The unreachable branch of the jump is removed
+        # The unreachable branch of the jump is removed, the jump
+        # becomes redundant and is replaced by a NOP (for the lineno)
 
         insts = [
             ('LOAD_CONST', 3, 11),
@@ -903,8 +904,7 @@ def test_conditional_jump_forward_const_condition(self):
         ]
         expected = [
             ('NOP', None, 11),
-            ('JUMP', lbl := self.Label(), 12),
-            lbl,
+            ('NOP', None, 12),
             ('LOAD_CONST', '3', 14)
         ]
         self.cfg_optimization_test(insts, expected, consts=list(range(5)))
diff --git a/Python/compile.c b/Python/compile.c
index 627f86a8ce9188..857fca440ac0a7 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7396,6 +7396,9 @@ mark_cold(basicblock *entryblock) {
     return 0;
 }
 
+static int
+remove_redundant_jumps(cfg_builder *g);
+
 static int
 push_cold_blocks_to_end(cfg_builder *g, int code_flags) {
     basicblock *entryblock = g->g_entryblock;
@@ -7465,6 +7468,12 @@ push_cold_blocks_to_end(cfg_builder *g, int code_flags) {
     }
     assert(b != NULL && b->b_next == NULL);
     b->b_next = cold_blocks;
+
+    if (cold_blocks != NULL) {
+        if (remove_redundant_jumps(g) < 0) {
+            return -1;
+        }
+    }
     return 0;
 }
 
@@ -8269,9 +8278,6 @@ trim_unused_consts(basicblock *entryblock, PyObject *consts);
 static int
 duplicate_exits_without_lineno(cfg_builder *g);
 
-static int
-extend_block(basicblock *bb);
-
 static int *
 build_cellfixedoffsets(struct compiler *c)
 {
@@ -8476,6 +8482,21 @@ propagate_line_numbers(basicblock *entryblock);
 static void
 eliminate_empty_basic_blocks(cfg_builder *g);
 
+#ifndef NDEBUG
+static bool
+no_redundant_jumps(cfg_builder *g) {
+    for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
+        struct instr *last = basicblock_last_instr(b);
+        if (last != NULL) {
+            if (last->i_opcode == JUMP || last->i_opcode == JUMP_NO_INTERRUPT) {
+                assert(last->i_target != b->b_next);
+                return false;
+            }
+        }
+    }
+    return true;
+}
+#endif
 
 static int
 remove_redundant_jumps(cfg_builder *g) {
@@ -8592,8 +8613,8 @@ assemble(struct compiler *c, int addNone)
     if (trim_unused_consts(g->g_entryblock, consts)) {
         goto error;
     }
-    if (duplicate_exits_without_lineno(g)) {
-        return NULL;
+    if (duplicate_exits_without_lineno(g) < 0) {
+        goto error;
     }
     propagate_line_numbers(g->g_entryblock);
     guarantee_lineno_for_exits(g->g_entryblock, c->u->u_firstlineno);
@@ -8612,10 +8633,6 @@ assemble(struct compiler *c, int addNone)
     if (push_cold_blocks_to_end(g, code_flags) < 0) {
         goto error;
     }
-
-    if (remove_redundant_jumps(g) < 0) {
-        goto error;
-    }
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
         clean_basic_block(b);
     }
@@ -8627,6 +8644,8 @@ assemble(struct compiler *c, int addNone)
         goto error;
     }
 
+    assert(no_redundant_jumps(g));
+
     /* Can't modify the bytecode after computing jump offsets. */
     assemble_jump_offsets(g->g_entryblock);
 
@@ -9488,7 +9507,7 @@ optimize_cfg(cfg_builder *g, PyObject *consts, PyObject *const_cache)
         }
     }
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
-        if (extend_block(b)) {
+        if (extend_block(b) < 0) {
             return -1;
         }
     }
@@ -9500,7 +9519,7 @@ optimize_cfg(cfg_builder *g, PyObject *consts, PyObject *const_cache)
         assert(b->b_predecessors == 0);
     }
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
-        if (extend_block(b)) {
+        if (extend_block(b) < 0) {
             return -1;
         }
     }
@@ -9517,6 +9536,9 @@ optimize_cfg(cfg_builder *g, PyObject *consts, PyObject *const_cache)
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
         clean_basic_block(b);
     }
+    if (remove_redundant_jumps(g) < 0) {
+        return -1;
+    }
     return 0;
 }