bpo-29469: Optimize literal lists and sets iterating on the AST level. (#4866)

Python/ast_opt.c

@@ -205,7 +205,7 @@ fold_binop(expr_ty node, PyArena *arena)
 }
 
 static PyObject*
-make_const_tuple(asdl_seq *elts, int make_set)
+make_const_tuple(asdl_seq *elts)
 {
     for (int i = 0; i < asdl_seq_LEN(elts); i++) {
         expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
@@ -225,11 +225,6 @@ make_const_tuple(asdl_seq *elts, int make_set)
         Py_INCREF(v);
         PyTuple_SET_ITEM(newval, i, v);
     }
-
-    /* Need to create frozen_set instead. */
-    if (make_set) {
-        Py_SETREF(newval, PyFrozenSet_New(newval));
-    }
     return newval;
 }
 
@@ -241,7 +236,7 @@ fold_tuple(expr_ty node, PyArena *arena)
     if (node->v.Tuple.ctx != Load)
         return 1;
 
-    newval = make_const_tuple(node->v.Tuple.elts, 0);
+    newval = make_const_tuple(node->v.Tuple.elts);
     return make_const(node, newval, arena);
 }
 
@@ -268,38 +263,48 @@ fold_subscr(expr_ty node, PyArena *arena)
     return make_const(node, newval, arena);
 }
 
+/* Change literal list or set of constants into constant
+   tuple or frozenset respectively.
+   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)
+{
+    PyObject *newval;
+    if (arg->kind == List_kind) {
+        newval = make_const_tuple(arg->v.List.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)
 {
     asdl_int_seq *ops;
     asdl_seq *args;
-    PyObject *newval;
     int i;
 
     ops = node->v.Compare.ops;
     args = node->v.Compare.comparators;
     /* TODO: optimize cases with literal arguments. */
-    for (i = 0; i < asdl_seq_LEN(ops); i++) {
-        int op;
-        expr_ty arg;
-        asdl_seq *elts;
-
-        op = asdl_seq_GET(ops, i);
-        arg = (expr_ty)asdl_seq_GET(args, i);
     /* Change literal list or set in 'in' or 'not in' into
        tuple or frozenset respectively. */
-        /* TODO: do the same when list or set is used as iterable
-           in for loop and comprehensions? */
-        if (op != In && op != NotIn)
-            continue;
-        if (arg->kind == List_kind)
-            elts = arg->v.List.elts;
-        else if (arg->kind == Set_kind)
-            elts = arg->v.Set.elts;
-        else continue;
-
-        newval = make_const_tuple(elts, arg->kind == Set_kind);
-        make_const(arg, newval, arena);
+    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)) {
+            return 0;
+        }
     }
     return 1;
 }
@@ -497,6 +502,8 @@ astfold_comprehension(comprehension_ty node_, PyArena* ctx_)
     CALL(astfold_expr, expr_ty, node_->target);
     CALL(astfold_expr, expr_ty, node_->iter);
     CALL_SEQ(astfold_expr, expr_ty, node_->ifs);
+
+    CALL(fold_iter, expr_ty, node_->iter);
     return 1;
 }
 
@@ -565,6 +572,8 @@ astfold_stmt(stmt_ty node_, PyArena* ctx_)
         CALL(astfold_expr, expr_ty, node_->v.For.iter);
         CALL_SEQ(astfold_stmt, stmt_ty, node_->v.For.body);
         CALL_SEQ(astfold_stmt, stmt_ty, 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);

Python/peephole.c

@@ -173,8 +173,6 @@ copy_op_arg(_Py_CODEUNIT *codestr, Py_ssize_t i, unsigned char op,
    new constant (c1, c2, ... cn) can be appended.
    Called with codestr pointing to the first LOAD_CONST.
    Bails out with no change if one or more of the LOAD_CONSTs is missing.
-   Also works for BUILD_LIST and BUILT_SET when followed by an "in" or "not in"
-   test; for BUILD_SET it assembles a frozenset rather than a tuple.
 */
 static Py_ssize_t
 fold_tuple_on_constants(_Py_CODEUNIT *codestr, Py_ssize_t c_start,
@@ -198,15 +196,6 @@ fold_tuple_on_constants(_Py_CODEUNIT *codestr, Py_ssize_t c_start,
         PyTuple_SET_ITEM(newconst, i, constant);
     }
 
-    /* If it's a BUILD_SET, use the PyTuple we just built to create a
-       PyFrozenSet, and use that as the constant instead: */
-    if (opcode == BUILD_SET) {
-        Py_SETREF(newconst, PyFrozenSet_New(newconst));
-        if (newconst == NULL) {
-            return -1;
-        }
-    }
-
     /* Append folded constant onto consts */
     len_consts = PyList_GET_SIZE(consts);
     if (PyList_Append(consts, newconst)) {
@@ -358,24 +347,15 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
                 CONST_STACK_POP(1);
                 break;
 
-                /* Try to fold tuples of constants (includes a case for lists
-                   and sets which are only used for "in" and "not in" tests).
+                /* Try to fold tuples of constants.
                    Skip over BUILD_SEQN 1 UNPACK_SEQN 1.
                    Replace BUILD_SEQN 2 UNPACK_SEQN 2 with ROT2.
                    Replace BUILD_SEQN 3 UNPACK_SEQN 3 with ROT3 ROT2. */
             case BUILD_TUPLE:
-            case BUILD_LIST:
-            case BUILD_SET:
                 j = get_arg(codestr, i);
                 if (j > 0 && CONST_STACK_LEN() >= j) {
                     h = lastn_const_start(codestr, op_start, j);
-                    if ((opcode == BUILD_TUPLE &&
-                          ISBASICBLOCK(blocks, h, op_start)) ||
-                         ((opcode == BUILD_LIST || opcode == BUILD_SET) &&
-                          ((nextop==COMPARE_OP &&
-                          (_Py_OPARG(codestr[nexti]) == PyCmp_IN ||
-                           _Py_OPARG(codestr[nexti]) == PyCmp_NOT_IN)) ||
-                          nextop == GET_ITER) && ISBASICBLOCK(blocks, h, i + 1))) {
+                    if (ISBASICBLOCK(blocks, h, op_start)) {
                         h = fold_tuple_on_constants(codestr, h, i + 1, opcode,
                                                     consts, CONST_STACK_LASTN(j), j);
                         if (h >= 0) {
@@ -387,8 +367,7 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
                 }
                 if (nextop != UNPACK_SEQUENCE  ||
                     !ISBASICBLOCK(blocks, op_start, i + 1) ||
-                    j != get_arg(codestr, nexti) ||
-                    opcode == BUILD_SET)
+                    j != get_arg(codestr, nexti))
                     break;
                 if (j < 2) {
                     fill_nops(codestr, op_start, nexti + 1);