gh-126835: Move constant tuple folding from ast_opt to CFG (#130769)

3 files changed, 111 insertions, 42 deletions

diff --git a/Python/ast_opt.c b/Python/ast_opt.c
index 4a191e919e4..8ee322fdd15 100644
--- a/Python/ast_opt.c
+++ b/Python/ast_opt.c
@@ -390,44 +390,6 @@ fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
     return 1;
 }
 
-static PyObject*
-make_const_tuple(asdl_expr_seq *elts)
-{
-    for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
-        expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
-        if (e->kind != Constant_kind) {
-            return NULL;
-        }
-    }
-
-    PyObject *newval = PyTuple_New(asdl_seq_LEN(elts));
-    if (newval == NULL) {
-        return NULL;
-    }
-
-    for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
-        expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
-        PyObject *v = e->v.Constant.value;
-        PyTuple_SET_ITEM(newval, i, Py_NewRef(v));
-    }
-    return newval;
-}
-
-static int
-fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
-{
-    if (state->syntax_check_only) {
-        return 1;
-    }
-    PyObject *newval;
-
-    if (node->v.Tuple.ctx != Load)
-        return 1;
-
-    newval = make_const_tuple(node->v.Tuple.elts);
-    return make_const(node, newval, arena);
-}
-
 static int astfold_mod(mod_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
 static int astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
 static int astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
@@ -620,7 +582,6 @@ astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
         break;
     case Tuple_kind:
         CALL_SEQ(astfold_expr, expr, node_->v.Tuple.elts);
-        CALL(fold_tuple, expr_ty, node_);
         break;
     case Name_kind:
         if (state->syntax_check_only) {
diff --git a/Python/codegen.c b/Python/codegen.c
index e1f647451f7..44d23ca6d93 100644
--- a/Python/codegen.c
+++ b/Python/codegen.c
@@ -1688,11 +1688,26 @@ codegen_typealias(compiler *c, stmt_ty s)
     return SUCCESS;
 }
 
+static bool
+is_const_tuple(asdl_expr_seq *elts)
+{
+    for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
+        expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
+        if (e->kind != Constant_kind) {
+            return false;
+        }
+    }
+    return true;
+}
+
 /* Return false if the expression is a constant value except named singletons.
    Return true otherwise. */
 static bool
 check_is_arg(expr_ty e)
 {
+    if (e->kind == Tuple_kind) {
+        return !is_const_tuple(e->v.Tuple.elts);
+    }
     if (e->kind != Constant_kind) {
         return true;
     }
diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index 8154129090b..cf3e74005ce 100644
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -1482,6 +1482,95 @@ fold_tuple_of_constants(basicblock *bb, int i, PyObject *consts, PyObject *const
     return instr_make_load_const(instr, const_tuple, consts, const_cache);
 }
 
+/* Replace:
+    BUILD_LIST 0
+    LOAD_CONST c1
+    LIST_APPEND 1
+    LOAD_CONST c2
+    LIST_APPEND 1
+    ...
+    LOAD_CONST cN
+    LIST_APPEND 1
+    CALL_INTRINSIC_1 INTRINSIC_LIST_TO_TUPLE
+   with:
+    LOAD_CONST (c1, c2, ... cN)
+*/
+static int
+fold_constant_intrinsic_list_to_tuple(basicblock *bb, int i,
+                                      PyObject *consts, PyObject *const_cache)
+{
+    assert(PyDict_CheckExact(const_cache));
+    assert(PyList_CheckExact(consts));
+    assert(i >= 0);
+    assert(i < bb->b_iused);
+
+    cfg_instr *intrinsic = &bb->b_instr[i];
+    assert(intrinsic->i_opcode == CALL_INTRINSIC_1);
+    assert(intrinsic->i_oparg == INTRINSIC_LIST_TO_TUPLE);
+
+    int consts_found = 0;
+    bool expect_append = true;
+
+    for (int pos = i - 1; pos >= 0; pos--) {
+        cfg_instr *instr = &bb->b_instr[pos];
+        int opcode = instr->i_opcode;
+        int oparg = instr->i_oparg;
+
+        if (opcode == NOP) {
+            continue;
+        }
+
+        if (opcode == BUILD_LIST && oparg == 0) {
+            if (!expect_append) {
+                /* Not a sequence start. */
+                return SUCCESS;
+            }
+
+            /* Sequence start, we are done. */
+            PyObject *newconst = PyTuple_New((Py_ssize_t)consts_found);
+            if (newconst == NULL) {
+                return ERROR;
+            }
+
+            for (int newpos = i - 1; newpos >= pos; newpos--) {
+                instr = &bb->b_instr[newpos];
+                if (instr->i_opcode == NOP) {
+                    continue;
+                }
+                if (loads_const(instr->i_opcode)) {
+                    PyObject *constant = get_const_value(instr->i_opcode, instr->i_oparg, consts);
+                    if (constant == NULL) {
+                        Py_DECREF(newconst);
+                        return ERROR;
+                    }
+                    assert(consts_found > 0);
+                    PyTuple_SET_ITEM(newconst, --consts_found, constant);
+                }
+                nop_out(&instr, 1);
+            }
+            assert(consts_found == 0);
+            return instr_make_load_const(intrinsic, newconst, consts, const_cache);
+        }
+
+        if (expect_append) {
+            if (opcode != LIST_APPEND || oparg != 1) {
+                return SUCCESS;
+            }
+        }
+        else {
+            if (!loads_const(opcode)) {
+                return SUCCESS;
+            }
+            consts_found++;
+        }
+
+        expect_append = !expect_append;
+    }
+
+    /* Did not find sequence start. */
+    return SUCCESS;
+}
+
 #define MIN_CONST_SEQUENCE_SIZE 3
 /*
 Optimize lists and sets for:
@@ -2378,9 +2467,13 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
                 RETURN_IF_ERROR(fold_const_unaryop(bb, i, consts, const_cache));
                 break;
             case CALL_INTRINSIC_1:
-                // for _ in (*foo, *bar) -> for _ in [*foo, *bar]
-                if (oparg == INTRINSIC_LIST_TO_TUPLE && nextop == GET_ITER) {
-                    INSTR_SET_OP0(inst, NOP);
+                if (oparg == INTRINSIC_LIST_TO_TUPLE) {
+                    if (nextop == GET_ITER) {
+                        INSTR_SET_OP0(inst, NOP);
+                    }
+                    else {
+                        RETURN_IF_ERROR(fold_constant_intrinsic_list_to_tuple(bb, i, consts, const_cache));
+                    }
                 }
                 else if (oparg == INTRINSIC_UNARY_POSITIVE) {
                     RETURN_IF_ERROR(fold_const_unaryop(bb, i, consts, const_cache));