gh-92782: unify the style of CFG traversal algorithms in the compiler (GH-92784)

1 files changed, 45 insertions, 38 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 51ef8fd..fdf2e5b 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7061,7 +7061,6 @@ struct assembler {
     PyObject *a_except_table;  /* bytes containing exception table */
     basicblock *a_entry;
     int a_offset;              /* offset into bytecode */
-    int a_nblocks;             /* number of reachable blocks */
     int a_except_table_off;    /* offset into exception table */
     int a_prevlineno;     /* lineno of last emitted line in line table */
     int a_prev_end_lineno; /* end_lineno of last emitted line in line table */
@@ -7074,6 +7073,20 @@ struct assembler {
     int a_location_off;    /* offset of last written location info frame */
 };
 
+static basicblock**
+make_cfg_traversal_stack(basicblock *entry) {
+    int nblocks = 0;
+    for (basicblock *b = entry; b != NULL; b = b->b_next) {
+        b->b_visited = 0;
+        nblocks++;
+    }
+    basicblock **stack = (basicblock **)PyMem_Malloc(sizeof(basicblock *) * nblocks);
+    if (!stack) {
+        PyErr_NoMemory();
+    }
+    return stack;
+}
+
 Py_LOCAL_INLINE(void)
 stackdepth_push(basicblock ***sp, basicblock *b, int depth)
 {
@@ -7089,31 +7102,26 @@ stackdepth_push(basicblock ***sp, basicblock *b, int depth)
  * cycles in the flow graph have no net effect on the stack depth.
  */
 static int
-stackdepth(struct compiler *c)
+stackdepth(struct compiler *c, basicblock *entry)
 {
-    basicblock *b, *entryblock = NULL;
-    basicblock **stack, **sp;
-    int nblocks = 0, maxdepth = 0;
-    for (b = c->u->u_blocks; b != NULL; b = b->b_list) {
+    for (basicblock *b = entry; b != NULL; b = b->b_next) {
         b->b_startdepth = INT_MIN;
-        entryblock = b;
-        nblocks++;
     }
-    assert(entryblock!= NULL);
-    stack = (basicblock **)PyObject_Malloc(sizeof(basicblock *) * nblocks);
+    basicblock **stack = make_cfg_traversal_stack(entry);
     if (!stack) {
-        PyErr_NoMemory();
         return -1;
     }
 
-    sp = stack;
+    int maxdepth = 0;
+    basicblock **sp = stack;
     if (c->u->u_ste->ste_generator || c->u->u_ste->ste_coroutine) {
-        stackdepth_push(&sp, entryblock, 1);
+        stackdepth_push(&sp, entry, 1);
     } else {
-        stackdepth_push(&sp, entryblock, 0);
+        stackdepth_push(&sp, entry, 0);
     }
+
     while (sp != stack) {
-        b = *--sp;
+        basicblock *b = *--sp;
         int depth = b->b_startdepth;
         assert(depth >= 0);
         basicblock *next = b->b_next;
@@ -7159,7 +7167,7 @@ stackdepth(struct compiler *c)
             stackdepth_push(&sp, next, depth);
         }
     }
-    PyObject_Free(stack);
+    PyMem_Free(stack);
     return maxdepth;
 }
 
@@ -7264,14 +7272,8 @@ copy_except_stack(ExceptStack *stack) {
 
 static int
 label_exception_targets(basicblock *entry) {
-    int nblocks = 0;
-    for (basicblock *b = entry; b != NULL; b = b->b_next) {
-        b->b_visited = 0;
-        nblocks++;
-    }
-    basicblock **todo_stack = PyMem_Malloc(sizeof(basicblock *)*nblocks);
+    basicblock **todo_stack = make_cfg_traversal_stack(entry);
     if (todo_stack == NULL) {
-        PyErr_NoMemory();
         return -1;
     }
     ExceptStack *except_stack = make_except_stack();
@@ -8051,7 +8053,7 @@ static int
 optimize_cfg(struct compiler *c, struct assembler *a, PyObject *consts);
 
 static int
-trim_unused_consts(struct compiler *c, struct assembler *a, PyObject *consts);
+trim_unused_consts(struct assembler *a, PyObject *consts);
 
 /* Duplicates exit BBs, so that line numbers can be propagated to them */
 static int
@@ -8347,7 +8349,6 @@ assemble(struct compiler *c, int addNone)
     if (!assemble_init(&a, nblocks, c->u->u_firstlineno))
         goto error;
     a.a_entry = entryblock;
-    a.a_nblocks = nblocks;
 
     int numdropped = fix_cell_offsets(c, entryblock, cellfixedoffsets);
     PyMem_Free(cellfixedoffsets);  // At this point we're done with it.
@@ -8368,12 +8369,12 @@ assemble(struct compiler *c, int addNone)
     if (duplicate_exits_without_lineno(c)) {
         return NULL;
     }
-    if (trim_unused_consts(c, &a, consts)) {
+    if (trim_unused_consts(&a, consts)) {
         goto error;
     }
     propagate_line_numbers(&a);
     guarantee_lineno_for_exits(&a, c->u->u_firstlineno);
-    int maxdepth = stackdepth(c);
+    int maxdepth = stackdepth(c, entryblock);
     if (maxdepth < 0) {
         goto error;
     }
@@ -9081,17 +9082,19 @@ normalize_basic_block(basicblock *bb) {
 
 static int
 mark_reachable(struct assembler *a) {
-    basicblock **stack, **sp;
-    sp = stack = (basicblock **)PyObject_Malloc(sizeof(basicblock *) * a->a_nblocks);
+    basicblock **stack = make_cfg_traversal_stack(a->a_entry);
     if (stack == NULL) {
         return -1;
     }
+    basicblock **sp = stack;
     a->a_entry->b_predecessors = 1;
     *sp++ = a->a_entry;
     while (sp > stack) {
         basicblock *b = *(--sp);
+        b->b_visited = 1;
         if (b->b_next && !b->b_nofallthrough) {
-            if (b->b_next->b_predecessors == 0) {
+            if (!b->b_next->b_visited) {
+                assert(b->b_next->b_predecessors == 0);
                 *sp++ = b->b_next;
             }
             b->b_next->b_predecessors++;
@@ -9101,14 +9104,15 @@ mark_reachable(struct assembler *a) {
             struct instr *instr = &b->b_instr[i];
             if (is_jump(instr) || is_block_push(instr)) {
                 target = instr->i_target;
-                if (target->b_predecessors == 0) {
+                if (!target->b_visited) {
+                    assert(target->b_predecessors == 0 || target == b->b_next);
                     *sp++ = target;
                 }
                 target->b_predecessors++;
             }
         }
     }
-    PyObject_Free(stack);
+    PyMem_Free(stack);
     return 0;
 }
 
@@ -9128,12 +9132,15 @@ eliminate_empty_basic_blocks(basicblock *entry) {
         if (b->b_iused == 0) {
             continue;
         }
-        if (is_jump(&b->b_instr[b->b_iused-1])) {
-            basicblock *target = b->b_instr[b->b_iused-1].i_target;
-            while (target->b_iused == 0) {
-                target = target->b_next;
+        for (int i = 0; i < b->b_iused; i++) {
+            struct instr *instr = &b->b_instr[i];
+            if (is_jump(instr) || is_block_push(instr)) {
+                basicblock *target = instr->i_target;
+                while (target->b_iused == 0) {
+                    target = target->b_next;
+                }
+                instr->i_target = target;
             }
-            b->b_instr[b->b_iused-1].i_target = target;
         }
     }
 }
@@ -9253,7 +9260,7 @@ optimize_cfg(struct compiler *c, struct assembler *a, PyObject *consts)
 
 // Remove trailing unused constants.
 static int
-trim_unused_consts(struct compiler *c, struct assembler *a, PyObject *consts)
+trim_unused_consts(struct assembler *a, PyObject *consts)
 {
     assert(PyList_CheckExact(consts));