gh-126835: Move const folding of lists & sets from ast_opt.c to flowgraph.c (#130032)

2 files changed, 34 insertions, 79 deletions

diff --git a/Python/ast_opt.c b/Python/ast_opt.c
index 78d84002d59..0b58e8cd2a2 100644
--- a/Python/ast_opt.c
+++ b/Python/ast_opt.c
@@ -567,62 +567,6 @@ fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
     return make_const(node, newval, arena);
 }
 
-/* Change literal list or set of constants into constant
-   tuple or frozenset respectively.  Change literal list of
-   non-constants into tuple.
-   Used for right operand of "in" and "not in" tests and for iterable
-   in "for" loop and comprehensions.
-*/
-static int
-fold_iter(expr_ty arg, PyArena *arena, _PyASTOptimizeState *state)
-{
-    PyObject *newval;
-    if (arg->kind == List_kind) {
-        /* First change a list into tuple. */
-        asdl_expr_seq *elts = arg->v.List.elts;
-        if (has_starred(elts)) {
-            return 1;
-        }
-        expr_context_ty ctx = arg->v.List.ctx;
-        arg->kind = Tuple_kind;
-        arg->v.Tuple.elts = elts;
-        arg->v.Tuple.ctx = ctx;
-        /* Try to create a constant tuple. */
-        newval = make_const_tuple(elts);
-    }
-    else if (arg->kind == Set_kind) {
-        newval = make_const_tuple(arg->v.Set.elts);
-        if (newval) {
-            Py_SETREF(newval, PyFrozenSet_New(newval));
-        }
-    }
-    else {
-        return 1;
-    }
-    return make_const(arg, newval, arena);
-}
-
-static int
-fold_compare(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
-{
-    asdl_int_seq *ops;
-    asdl_expr_seq *args;
-    Py_ssize_t i;
-
-    ops = node->v.Compare.ops;
-    args = node->v.Compare.comparators;
-    /* Change literal list or set in 'in' or 'not in' into
-       tuple or frozenset respectively. */
-    i = asdl_seq_LEN(ops) - 1;
-    int op = asdl_seq_GET(ops, i);
-    if (op == In || op == NotIn) {
-        if (!fold_iter((expr_ty)asdl_seq_GET(args, i), arena, state)) {
-            return 0;
-        }
-    }
-    return 1;
-}
-
 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);
@@ -783,7 +727,6 @@ astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
     case Compare_kind:
         CALL(astfold_expr, expr_ty, node_->v.Compare.left);
         CALL_SEQ(astfold_expr, expr, node_->v.Compare.comparators);
-        CALL(fold_compare, expr_ty, node_);
         break;
     case Call_kind:
         CALL(astfold_expr, expr_ty, node_->v.Call.func);
@@ -852,8 +795,6 @@ astfold_comprehension(comprehension_ty node_, PyArena *ctx_, _PyASTOptimizeState
     CALL(astfold_expr, expr_ty, node_->target);
     CALL(astfold_expr, expr_ty, node_->iter);
     CALL_SEQ(astfold_expr, expr, node_->ifs);
-
-    CALL(fold_iter, expr_ty, node_->iter);
     return 1;
 }
 
@@ -940,8 +881,6 @@ astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
         CALL(astfold_expr, expr_ty, node_->v.For.iter);
         CALL_SEQ(astfold_stmt, stmt, node_->v.For.body);
         CALL_SEQ(astfold_stmt, stmt, node_->v.For.orelse);
-
-        CALL(fold_iter, expr_ty, node_->v.For.iter);
         break;
     case AsyncFor_kind:
         CALL(astfold_expr, expr_ty, node_->v.AsyncFor.target);
diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index 308cdac3c44..f9a7c6fe210 100644
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -1428,31 +1428,41 @@ fold_tuple_of_constants(basicblock *bb, int n, PyObject *consts, PyObject *const
 }
 
 #define MIN_CONST_SEQUENCE_SIZE 3
-/* Replace LOAD_CONST c1, LOAD_CONST c2 ... LOAD_CONST cN, BUILD_LIST N
-   with BUILD_LIST 0, LOAD_CONST (c1, c2, ... cN), LIST_EXTEND 1,
-   or BUILD_SET & SET_UPDATE respectively.
+/*
+Optimize lists and sets for:
+    1. "for" loop, comprehension or "in"/"not in" tests:
+           Change literal list or set of constants into constant
+           tuple or frozenset respectively. Change list of
+           non-constants into tuple.
+    2. Constant literal lists/set with length >= MIN_CONST_SEQUENCE_SIZE:
+           Replace LOAD_CONST c1, LOAD_CONST c2 ... LOAD_CONST cN, BUILD_LIST N
+           with BUILD_LIST 0, LOAD_CONST (c1, c2, ... cN), LIST_EXTEND 1,
+           or BUILD_SET & SET_UPDATE respectively.
 */
 static int
-optimize_if_const_list_or_set(basicblock *bb, int n, PyObject *consts, PyObject *const_cache)
+optimize_lists_and_sets(basicblock *bb, int i, int nextop,
+                        PyObject *consts, PyObject *const_cache)
 {
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
-    cfg_instr *instr = &bb->b_instr[n];
+    cfg_instr *instr = &bb->b_instr[i];
     assert(instr->i_opcode == BUILD_LIST || instr->i_opcode == BUILD_SET);
+    bool contains_or_iter = nextop == GET_ITER || nextop == CONTAINS_OP;
     int seq_size = instr->i_oparg;
-    if (seq_size < MIN_CONST_SEQUENCE_SIZE) {
+    if (seq_size < MIN_CONST_SEQUENCE_SIZE && !contains_or_iter) {
         return SUCCESS;
     }
     PyObject *newconst;
-    RETURN_IF_ERROR(get_constant_sequence(bb, n-1, seq_size, consts, &newconst));
-    if (newconst == NULL) {
-        /* not a const sequence */
+    RETURN_IF_ERROR(get_constant_sequence(bb, i-1, seq_size, consts, &newconst));
+    if (newconst == NULL) {    /* not a const sequence */
+        if (contains_or_iter && instr->i_opcode == BUILD_LIST) {
+            /* iterate over a tuple instead of list */
+            INSTR_SET_OP1(instr, BUILD_TUPLE, instr->i_oparg);
+        }
         return SUCCESS;
     }
     assert(PyTuple_CheckExact(newconst) && PyTuple_GET_SIZE(newconst) == seq_size);
-    int build = instr->i_opcode;
-    int extend = build == BUILD_LIST ? LIST_EXTEND : SET_UPDATE;
-    if (build == BUILD_SET) {
+    if (instr->i_opcode == BUILD_SET) {
         PyObject *frozenset = PyFrozenSet_New(newconst);
         if (frozenset == NULL) {
             Py_DECREF(newconst);
@@ -1462,11 +1472,17 @@ optimize_if_const_list_or_set(basicblock *bb, int n, PyObject *consts, PyObject
     }
     int index = add_const(newconst, consts, const_cache);
     RETURN_IF_ERROR(index);
-    nop_out(bb, n-1, seq_size);
-    assert(n >= 2);
-    INSTR_SET_OP1(&bb->b_instr[n-2], build, 0);
-    INSTR_SET_OP1(&bb->b_instr[n-1], LOAD_CONST, index);
-    INSTR_SET_OP1(&bb->b_instr[n], extend, 1);
+    nop_out(bb, i-1, seq_size);
+    if (contains_or_iter) {
+        INSTR_SET_OP1(instr, LOAD_CONST, index);
+    }
+    else {
+        assert(i >= 2);
+        assert(instr->i_opcode == BUILD_LIST || instr->i_opcode == BUILD_SET);
+        INSTR_SET_OP1(&bb->b_instr[i-2], instr->i_opcode, 0);
+        INSTR_SET_OP1(&bb->b_instr[i-1], LOAD_CONST, index);
+        INSTR_SET_OP1(&bb->b_instr[i], instr->i_opcode == BUILD_LIST ? LIST_EXTEND : SET_UPDATE, 1);
+    }
     return SUCCESS;
 }
 
@@ -1923,7 +1939,7 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
                 break;
             case BUILD_LIST:
             case BUILD_SET:
-                RETURN_IF_ERROR(optimize_if_const_list_or_set(bb, i, consts, const_cache));
+                RETURN_IF_ERROR(optimize_lists_and_sets(bb, i, nextop, consts, const_cache));
                 break;
             case POP_JUMP_IF_NOT_NONE:
             case POP_JUMP_IF_NONE: