Fold set into frozenset in CFG

Author: Eclips4

Lib/test/test_ast/test_ast.py

@@ -3228,7 +3228,6 @@ def test_folding_comparator(self):
         operators = [("in", ast.In()), ("not in", ast.NotIn())]
         braces = [
             ("[", "]", ast.List, (1,)),
-            ("{", "}", ast.Set, frozenset({1})),
         ]
         for left, right, non_optimized_comparator, optimized_comparator in braces:
             for op, node in operators:
@@ -3246,7 +3245,6 @@ def test_folding_iter(self):
         code = "for _ in %s1%s: pass"
         braces = [
             ("[", "]", ast.List, (1,)),
-            ("{", "}", ast.Set, frozenset({1})),
         ]
 
         for left, right, ast_cls, optimized_iter in braces:

Python/ast_opt.c

@@ -575,8 +575,7 @@ fold_subscr(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
 }
 
 /* Change literal list or set of constants into constant
-   tuple or frozenset respectively.  Change literal list of
-   non-constants into tuple.
+   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.
 */
@@ -597,12 +596,6 @@ fold_iter(expr_ty arg, PyArena *arena, _PyASTOptimizeState *state)
         /* 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;
     }
@@ -618,8 +611,7 @@ fold_compare(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
 
     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. */
+    /* 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) {

Python/flowgraph.c

@@ -1336,6 +1336,18 @@ add_const(PyObject *newconst, PyObject *consts, PyObject *const_cache)
     return (int)index;
 }
 
+static int
+is_sequence_constant(cfg_instr *inst, int n)
+{
+    for (int i = 0; i < n; i++) {
+        if (!loads_const(inst[i].i_opcode)) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
+
 /* Replace LOAD_CONST c1, LOAD_CONST c2 ... LOAD_CONST cn, BUILD_TUPLE n
    with    LOAD_CONST (c1, c2, ... cn).
    The consts table must still be in list form so that the
@@ -1353,17 +1365,60 @@ fold_tuple_on_constants(PyObject *const_cache,
     assert(inst[n].i_opcode == BUILD_TUPLE);
     assert(inst[n].i_oparg == n);
 
+    if (!is_sequence_constant(inst, n)) {
+        return SUCCESS;
+    }
+
+    /* Buildup new tuple of constants */
+    PyObject *newconst = PyTuple_New(n);
+    if (newconst == NULL) {
+        return ERROR;
+    }
     for (int i = 0; i < n; i++) {
-        if (!loads_const(inst[i].i_opcode)) {
-            return SUCCESS;
+        int op = inst[i].i_opcode;
+        int arg = inst[i].i_oparg;
+        PyObject *constant = get_const_value(op, arg, consts);
+        if (constant == NULL) {
+            return ERROR;
         }
+        PyTuple_SET_ITEM(newconst, i, constant);
+    }
+    int index = add_const(newconst, consts, const_cache);
+    if (index < 0) {
+        return ERROR;
+    }
+    for (int i = 0; i < n; i++) {
+        INSTR_SET_OP0(&inst[i], NOP);
+    }
+    INSTR_SET_OP1(&inst[n], LOAD_CONST, index);
+    return SUCCESS;
+}
+
+
+// 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.
+static int
+fold_set_on_constants(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);
+
+    if (!is_sequence_constant(inst, n)) {
+        return SUCCESS;
     }
 
-    /* Buildup new tuple of constants */
     PyObject *newconst = PyTuple_New(n);
     if (newconst == NULL) {
         return ERROR;
     }
+
     for (int i = 0; i < n; i++) {
         int op = inst[i].i_opcode;
         int arg = inst[i].i_oparg;
@@ -1373,6 +1428,13 @@ fold_tuple_on_constants(PyObject *const_cache,
         }
         PyTuple_SET_ITEM(newconst, i, constant);
     }
+
+    PyObject *frozenset = PyFrozenSet_New(newconst);
+    if (frozenset == NULL) {
+        return ERROR;
+    }
+    Py_SETREF(newconst, frozenset);
+
     int index = add_const(newconst, consts, const_cache);
     if (index < 0) {
         return ERROR;
@@ -1751,6 +1813,13 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
                     }
                 }
                 break;
+            case BUILD_SET:
+                if (nextop == CONTAINS_OP || nextop == GET_ITER) {
+                    if (fold_set_on_constants(const_cache, inst-oparg, oparg, consts)) {
+                        goto error;
+                    }
+                }
+                break;
             case POP_JUMP_IF_NOT_NONE:
             case POP_JUMP_IF_NONE:
                 switch (target->i_opcode) {