gh-148817: Fold long lists/sets of constant elements into constant tuples/frozensets (GH-149016)

Fold long lists/sets of constant elements into constant tuples/frozensets.

(cherry picked from commit 084230e)

Co-authored-by: Kirill Podoprigora <kirill.bast9@mail.ru>
Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>

Author: 2 people

.github/CODEOWNERS

@@ -227,7 +227,7 @@ Tools/cases_generator/        @markshannon
 Python/assemble.c             @markshannon @iritkatriel
 Python/codegen.c              @markshannon @iritkatriel
 Python/compile.c              @markshannon @iritkatriel
-Python/flowgraph.c            @markshannon @iritkatriel
+Python/flowgraph.c            @markshannon @iritkatriel @eclips4
 Python/instruction_sequence.c @iritkatriel
 Python/symtable.c             @JelleZijlstra @carljm
 

Lib/test/test_peepholer.py

@@ -2470,6 +2470,168 @@ def test_list_to_tuple_get_iter_is_safe(self):
         self.assertEqual(b, [3, 2, 1, 0])
         self.assertEqual(items, [])
 
+    def test_fold_constant_big_list_for_iter(self):
+        # for x in [c1, c2, ..., cN] (N > 30) should fold to LOAD_CONST tuple
+        consts = 35
+        before = (
+            [("BUILD_LIST", 0, 1)] +
+            [("LOAD_CONST", 0, 2), ("LIST_APPEND", 1, 3)] * consts +
+            [("GET_ITER", 0, 4),
+             top := self.Label(),
+             ("FOR_ITER", end := self.Label(), 5),
+             ("STORE_FAST", 0, 6),
+             ("JUMP", top, 7),
+             end,
+             ("END_FOR", None, 8),
+             ("POP_ITER", None, 9),
+             ("LOAD_CONST", 0, 10),
+             ("RETURN_VALUE", None, 11)]
+        )
+        after = [
+            ("LOAD_CONST", 1, 3),
+            ("GET_ITER", 0, 4),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 5),
+            ("STORE_FAST", 0, 6),
+            ("JUMP", top, 7),
+            end,
+            ("END_FOR", None, 8),
+            ("POP_ITER", None, 9),
+            ("LOAD_CONST", 0, 10),
+            ("RETURN_VALUE", None, 11),
+        ]
+        result_const = tuple(["test"] * consts)
+        self.cfg_optimization_test(before, after, consts=["test"],
+                                   expected_consts=["test", result_const])
+
+    def test_fold_constant_big_set_for_iter(self):
+        # for x in {c1, c2, ..., cN} (N > 30) should fold to LOAD_CONST frozenset
+        before = [
+            ("BUILD_SET", 0, 1),
+            ("LOAD_SMALL_INT", 1, 2), ("SET_ADD", 1, 3),
+            ("LOAD_SMALL_INT", 2, 4), ("SET_ADD", 1, 5),
+            ("LOAD_SMALL_INT", 3, 6), ("SET_ADD", 1, 7),
+            ("GET_ITER", 0, 8),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 9),
+            ("STORE_FAST", 0, 10),
+            ("JUMP", top, 11),
+            end,
+            ("END_FOR", None, 12),
+            ("POP_ITER", None, 13),
+            ("LOAD_CONST", 0, 14),
+            ("RETURN_VALUE", None, 15),
+        ]
+        after = [
+            ("LOAD_CONST", 1, 7),
+            ("GET_ITER", 0, 8),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 9),
+            ("STORE_FAST", 0, 10),
+            ("JUMP", top, 11),
+            end,
+            ("END_FOR", None, 12),
+            ("POP_ITER", None, 13),
+            ("LOAD_CONST", 0, 14),
+            ("RETURN_VALUE", None, 15),
+        ]
+        self.cfg_optimization_test(before, after, consts=["test"],
+                                   expected_consts=["test", frozenset({1, 2, 3})])
+
+    def test_fold_constant_list_to_tuple_for_iter(self):
+        INTRINSIC_LIST_TO_TUPLE = 6
+        before = [
+            ("BUILD_LIST", 0, 1),
+            ("LOAD_SMALL_INT", 1, 2), ("LIST_APPEND", 1, 3),
+            ("LOAD_SMALL_INT", 2, 4), ("LIST_APPEND", 1, 5),
+            ("LOAD_SMALL_INT", 3, 6), ("LIST_APPEND", 1, 7),
+            ("CALL_INTRINSIC_1", INTRINSIC_LIST_TO_TUPLE, 8),
+            ("GET_ITER", 0, 9),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 10),
+            ("STORE_FAST", 0, 11),
+            ("JUMP", top, 12),
+            end,
+            ("END_FOR", None, 13),
+            ("POP_ITER", None, 14),
+            ("LOAD_CONST", 0, 15),
+            ("RETURN_VALUE", None, 16),
+        ]
+        after = [
+            ("LOAD_CONST", 1, 8),
+            ("GET_ITER", 0, 9),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 10),
+            ("STORE_FAST", 0, 11),
+            ("JUMP", top, 12),
+            end,
+            ("END_FOR", None, 13),
+            ("POP_ITER", None, 14),
+            ("LOAD_CONST", 0, 15),
+            ("RETURN_VALUE", None, 16),
+        ]
+        self.cfg_optimization_test(before, after, consts=["test"],
+                                   expected_consts=["test", (1, 2, 3)])
+
+    def test_fold_constant_big_list_contains_op(self):
+        # x in [c1, c2, ..., cN] (N > 30) should fold to LOAD_CONST tuple
+        before = [
+            ("LOAD_FAST", 0, 1),
+            ("BUILD_LIST", 0, 2),
+            ("LOAD_SMALL_INT", 1, 3), ("LIST_APPEND", 1, 4),
+            ("LOAD_SMALL_INT", 2, 5), ("LIST_APPEND", 1, 6),
+            ("LOAD_SMALL_INT", 3, 7), ("LIST_APPEND", 1, 8),
+            ("CONTAINS_OP", 0, 9),
+            ("RETURN_VALUE", None, 10),
+        ]
+        after = [
+            ("LOAD_FAST_BORROW", 0, 1),
+            ("LOAD_CONST", 1, 8),
+            ("CONTAINS_OP", 0, 9),
+            ("RETURN_VALUE", None, 10),
+        ]
+        self.cfg_optimization_test(before, after, consts=[None],
+                                   expected_consts=[None, (1, 2, 3)])
+
+    def test_fold_constant_big_set_contains_op(self):
+        # x in {c1, c2, ..., cN} (N > 30) should fold to LOAD_CONST frozenset
+        before = [
+            ("LOAD_FAST", 0, 1),
+            ("BUILD_SET", 0, 2),
+            ("LOAD_SMALL_INT", 1, 3), ("SET_ADD", 1, 4),
+            ("LOAD_SMALL_INT", 2, 5), ("SET_ADD", 1, 6),
+            ("LOAD_SMALL_INT", 3, 7), ("SET_ADD", 1, 8),
+            ("CONTAINS_OP", 0, 9),
+            ("RETURN_VALUE", None, 10),
+        ]
+        after = [
+            ("LOAD_FAST_BORROW", 0, 1),
+            ("LOAD_CONST", 1, 8),
+            ("CONTAINS_OP", 0, 9),
+            ("RETURN_VALUE", None, 10),
+        ]
+        self.cfg_optimization_test(before, after, consts=[None],
+                                   expected_consts=[None, frozenset({1, 2, 3})])
+
+    def test_no_fold_big_list_for_iter_with_non_const(self):
+        same = [
+            ("BUILD_LIST", 0, 1),
+            ("LOAD_SMALL_INT", 1, 2), ("LIST_APPEND", 1, 3),
+            ("LOAD_FAST_BORROW", 0, 4), ("LIST_APPEND", 1, 5),
+            ("LOAD_SMALL_INT", 3, 6), ("LIST_APPEND", 1, 7),
+            ("GET_ITER", 0, 8),
+            top := self.Label(),
+            ("FOR_ITER", end := self.Label(), 9),
+            ("STORE_FAST", 1, 10),
+            ("JUMP", top, 11),
+            end,
+            ("END_FOR", None, 12),
+            ("POP_ITER", None, 13),
+            ("LOAD_CONST", 0, 14),
+            ("RETURN_VALUE", None, 15),
+        ]
+        self.cfg_optimization_test(same, same, consts=["test"])
+
 
 class OptimizeLoadFastTestCase(DirectCfgOptimizerTests):
     def make_bb(self, insts):

Misc/NEWS.d/next/Core_and_Builtins/2026-04-26-15-08-53.gh-issue-148817.cuN07H.rst

@@ -0,0 +1,5 @@
+Fold large constant list and set literals used as the iterable of a
+:keyword:`for` loop or ``in``/``not in`` test into a constant
+:class:`tuple` or :class:`frozenset`, restoring an optimization
+previously done by the AST optimizer that was lost when constant
+folding moved to the CFG.

Python/flowgraph.c

@@ -1563,34 +1563,48 @@ fold_tuple_of_constants(basicblock *bb, int i, PyObject *consts,
 }
 
 /* Replace:
-    BUILD_LIST 0
+    BUILD_LIST/BUILD_SET 0
     LOAD_CONST c1
-    LIST_APPEND 1
+    LIST_APPEND/SET_ADD 1
     LOAD_CONST c2
-    LIST_APPEND 1
+    LIST_APPEND/SET_ADD 1
     ...
     LOAD_CONST cN
-    LIST_APPEND 1
-    CALL_INTRINSIC_1 INTRINSIC_LIST_TO_TUPLE
+    LIST_APPEND/SET_ADD 1
+    [CALL_INTRINSIC_1 INTRINSIC_LIST_TO_TUPLE]   <-- optional
    with:
     LOAD_CONST (c1, c2, ... cN)
+   The instruction at `i` is either the LIST_TO_TUPLE intrinsic (so the
+   immediately preceding non-NOP instruction is expected to be a
+   LIST_APPEND, and only the BUILD_LIST/LIST_APPEND form is considered),
+   or the trailing LIST_APPEND or SET_ADD itself, in which case the
+   matching BUILD_LIST/BUILD_SET start is selected from its opcode, and
+   for sets the result is wrapped in a frozenset.
 */
 static int
-fold_constant_intrinsic_list_to_tuple(basicblock *bb, int i,
-                                      PyObject *consts, PyObject *const_cache,
-                                      _Py_hashtable_t *consts_index)
+fold_constant_seq_into_load_const(basicblock *bb, int i,
+                                  PyObject *consts, PyObject *const_cache,
+                                  _Py_hashtable_t *consts_index)
 {
     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);
-
+    cfg_instr *target = &bb->b_instr[i];
+    assert(target->i_opcode == LIST_APPEND || target->i_opcode == SET_ADD ||
+           (target->i_opcode == CALL_INTRINSIC_1 &&
+            target->i_oparg == INTRINSIC_LIST_TO_TUPLE));
+    bool expected_append = target->i_opcode == CALL_INTRINSIC_1;
+    int append_op = expected_append ? LIST_APPEND : target->i_opcode;
+    assert(append_op == LIST_APPEND || append_op == SET_ADD);
+    int build_op = append_op == LIST_APPEND ? BUILD_LIST : BUILD_SET;
     int consts_found = 0;
-    bool expect_append = true;
+    /* Walking backward from `i`, we expect LIST_APPEND/SET_ADD and
+       LOAD_CONST to alternate. If `i` is the trailing LIST_TO_TUPLE
+       intrinsic, the next instruction back is an APPEND. If `i` is the
+       trailing APPEND itself, the next instruction back is a LOAD_CONST. */
+    bool expect_append = expected_append;
 
     for (int pos = i - 1; pos >= 0; pos--) {
         cfg_instr *instr = &bb->b_instr[pos];
@@ -1601,7 +1615,7 @@ fold_constant_intrinsic_list_to_tuple(basicblock *bb, int i,
             continue;
         }
 
-        if (opcode == BUILD_LIST && oparg == 0) {
+        if (opcode == build_op && oparg == 0) {
             if (!expect_append) {
                 /* Not a sequence start. */
                 return SUCCESS;
@@ -1613,7 +1627,8 @@ fold_constant_intrinsic_list_to_tuple(basicblock *bb, int i,
                 return ERROR;
             }
 
-            for (int newpos = i - 1; newpos >= pos; newpos--) {
+            int newpos_start = expected_append ? i - 1 : i;
+            for (int newpos = newpos_start; newpos >= pos; newpos--) {
                 instr = &bb->b_instr[newpos];
                 if (instr->i_opcode == NOP) {
                     continue;
@@ -1630,11 +1645,20 @@ fold_constant_intrinsic_list_to_tuple(basicblock *bb, int i,
                 nop_out(&instr, 1);
             }
             assert(consts_found == 0);
-            return instr_make_load_const(intrinsic, newconst, consts, const_cache, consts_index);
+
+            if (build_op == BUILD_SET) {
+                PyObject *frozen = PyFrozenSet_New(newconst);
+                Py_DECREF(newconst);
+                if (frozen == NULL) {
+                    return ERROR;
+                }
+                newconst = frozen;
+            }
+            return instr_make_load_const(target, newconst, consts, const_cache, consts_index);
         }
 
         if (expect_append) {
-            if (opcode != LIST_APPEND || oparg != 1) {
+            if (opcode != append_op || oparg != 1) {
                 return SUCCESS;
             }
         }
@@ -2573,17 +2597,22 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts,
                 break;
             case CALL_INTRINSIC_1:
                 if (oparg == INTRINSIC_LIST_TO_TUPLE) {
-                    if (nextop == GET_ITER) {
+                    RETURN_IF_ERROR(fold_constant_seq_into_load_const(bb, i, consts, const_cache, consts_index));
+                    if (inst->i_opcode == CALL_INTRINSIC_1 && nextop == GET_ITER) {
                         INSTR_SET_OP0(inst, NOP);
                     }
-                    else {
-                        RETURN_IF_ERROR(fold_constant_intrinsic_list_to_tuple(bb, i, consts, const_cache, consts_index));
-                    }
                 }
                 else if (oparg == INTRINSIC_UNARY_POSITIVE) {
                     RETURN_IF_ERROR(fold_const_unaryop(bb, i, consts, const_cache, consts_index));
                 }
                 break;
+            case LIST_APPEND:
+            case SET_ADD:
+                if (oparg == 1 && (nextop == GET_ITER || nextop == CONTAINS_OP)) {
+                    RETURN_IF_ERROR(fold_constant_seq_into_load_const(
+                        bb, i, consts, const_cache, consts_index));
+                }
+                break;
             case BINARY_OP:
                 RETURN_IF_ERROR(fold_const_binop(bb, i, consts, const_cache, consts_index));
                 break;