for target optimization (set/list -> frozenset/tuple)

tomasr8 · Jan 3, 2025 · Jan 3, 2025 · Jan 16, 2025 · Jan 17, 2025 · Jan 21, 2025
commit 17dffaa2f36adaf3492a8bf6587994b054811c18
@@ -3223,44 +3223,6 @@ def test_folding_format(self):
 
         self.assert_ast(code, non_optimized_target, optimized_target)
 
-    def test_folding_comparator(self):
-        code = "1 %s %s1%s"
-        operators = [("in", ast.In()), ("not in", ast.NotIn())]
-        braces = [
-            ("[", "]", ast.List, (1,)),
-        ]
-        for left, right, non_optimized_comparator, optimized_comparator in braces:
-            for op, node in operators:
-                non_optimized_target = self.wrap_expr(ast.Compare(
-                    left=ast.Constant(1), ops=[node],
-                    comparators=[non_optimized_comparator(elts=[ast.Constant(1)])]
-                ))
-                optimized_target = self.wrap_expr(ast.Compare(
-                    left=ast.Constant(1), ops=[node],
-                    comparators=[ast.Constant(value=optimized_comparator)]
-                ))
-                self.assert_ast(code % (op, left, right), non_optimized_target, optimized_target)
-
-    def test_folding_iter(self):
-        code = "for _ in %s1%s: pass"
-        braces = [
-            ("[", "]", ast.List, (1,)),
-        ]
-
-        for left, right, ast_cls, optimized_iter in braces:
-            non_optimized_target = self.wrap_statement(ast.For(
- ctx=ast.Store()),
-                iter=ast_cls(elts=[ast.Constant(1)]),
-                body=[ast.Pass()]
-            ))
-            optimized_target = self.wrap_statement(ast.For(
- ctx=ast.Store()),
-                iter=ast.Constant(value=optimized_iter),
-                body=[ast.Pass()]
-            ))
-
-            self.assert_ast(code % (left, right), non_optimized_target, optimized_target)
-
     def test_folding_subscript(self):
         code = "'abcd'[0]"
 

diff --git a/Python/ast_opt.c b/Python/ast_opt.c
@@ -532,29 +532,6 @@ fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
     return make_const(node, newval, arena);
 }
 
-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_subscr(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
 {
@@ -574,54 +551,6 @@ fold_subscr(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
     return make_const(node, newval, arena);
 }
 
-/* Change literal list or set of constants into constant
-   tuple. 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 {
-        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. */
-    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);
@@ -782,7 +711,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,7 +780,6 @@ astfold_comprehension(comprehension_ty node_, PyArena *ctx_, _PyASTOptimizeState
     CALL(astfold_expr, expr_ty, node_->iter);
     CALL_SEQ(astfold_expr, expr, node_->ifs);
 
-    CALL(fold_iter, expr_ty, node_->iter);
     return 1;
 }
 
@@ -939,8 +866,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
@@ -1395,21 +1395,24 @@ fold_tuple_on_constants(PyObject *const_cache,
 }
 
 
-// Replaces const set with a frozenset.
-// This should be used only in situations where we 100% sure that
-// this set cannot be changed: where's constant set is a rhs in `for` loop
-// or it's a rhs in `in` operation.
+/* Replaces const set/list with a frozenset/tuple.
+   This should be used only in situations where we 100% sure that
+   this set cannot be changed: where's constant set/list is a rhs in `for` loop
+   or it's a rhs in `in` operation.
+*/
 static int
-fold_set_on_constants(PyObject *const_cache,
+fold_if_const_list_or_set(PyObject *const_cache,
                       cfg_instr *inst,
                       int n, PyObject *consts)
 {
     /* Pre-conditions */
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
-    assert(inst[n].i_opcode == BUILD_SET);
     assert(inst[n].i_oparg == n);
 
+    int build = inst[n].i_opcode;
+    assert(build == BUILD_LIST || build == BUILD_SET);
+
     if (!is_constant_sequence(inst, n)) {
         return SUCCESS;
     }
@@ -1428,13 +1431,13 @@ fold_set_on_constants(PyObject *const_cache,
         }
         PyTuple_SET_ITEM(newconst, i, constant);
     }
-
-    PyObject *frozenset = PyFrozenSet_New(newconst);
-    if (frozenset == NULL) {
-        return ERROR;
+    if (build == BUILD_SET) {
+        PyObject *frozenset = PyFrozenSet_New(newconst);
+        if (frozenset == NULL) {
+            return ERROR;
+        }
+        Py_SETREF(newconst, frozenset);
     }
-    Py_SETREF(newconst, frozenset);
-
     int index = add_const(newconst, consts, const_cache);
     if (index < 0) {
         return ERROR;
@@ -1866,9 +1869,10 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
             case BUILD_LIST:
             case BUILD_SET:
                 if (nextop == CONTAINS_OP || nextop == GET_ITER) {
-                    if (fold_set_on_constants(const_cache, inst-oparg, oparg, consts) < 0) {
+                    if (fold_if_const_list_or_set(const_cache, inst-oparg, oparg, consts) < 0) {
                         goto error;
                     }
+                    break;
                 }
                 if (i >= oparg) {
                     if (optimize_if_const_list_or_set(const_cache, inst-oparg, oparg, consts) < 0) {