bpo-34100: compile: Re-enable frozenset merging (GH-10760)

methane · vstinner · commit f7e4d3642fbb · 2018-11-28T16:58:46.000+01:00
This reverts commit 1005c84.
diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
@@ -615,6 +615,14 @@ def check_same_constant(const):
         self.check_constant(f1, Ellipsis)
         self.assertEqual(repr(f1()), repr(Ellipsis))
 
+        # Merge constants in tuple or frozenset
+        f1, f2 = lambda: "not a name", lambda: ("not a name",)
+        f3 = lambda x: x in {("not a name",)}
+        self.assertIs(f1.__code__.co_consts[1],
+                      f2.__code__.co_consts[1][0])
+        self.assertIs(next(iter(f3.__code__.co_consts[1])),
+                      f2.__code__.co_consts[1])
+
         # {0} is converted to a constant frozenset({0}) by the peephole
         # optimizer
         f1, f2 = lambda x: x in {0}, lambda x: x in {0}
diff --git a/Python/compile.c b/Python/compile.c
@@ -1208,14 +1208,18 @@ merge_consts_recursive(struct compiler *c, PyObject *o)
     // t is borrowed reference
     PyObject *t = PyDict_SetDefault(c->c_const_cache, key, key);
     if (t != key) {
+        // o is registered in c_const_cache.  Just use it.
         Py_INCREF(t);
         Py_DECREF(key);
         return t;
     }
 
+    // We registered o in c_const_cache.
+    // When o is a tuple or frozenset, we want to merge it's
+    // items too.
     if (PyTuple_CheckExact(o)) {
-        Py_ssize_t i, len = PyTuple_GET_SIZE(o);
-        for (i = 0; i < len; i++) {
+        Py_ssize_t len = PyTuple_GET_SIZE(o);
+        for (Py_ssize_t i = 0; i < len; i++) {
             PyObject *item = PyTuple_GET_ITEM(o, i);
             PyObject *u = merge_consts_recursive(c, item);
             if (u == NULL) {
@@ -1240,6 +1244,57 @@ merge_consts_recursive(struct compiler *c, PyObject *o)
             Py_DECREF(u);
         }
     }
+    else if (PyFrozenSet_CheckExact(o)) {
+        // *key* is tuple. And it's first item is frozenset of
+        // constant keys.
+        // See _PyCode_ConstantKey() for detail.
+        assert(PyTuple_CheckExact(key));
+        assert(PyTuple_GET_SIZE(key) == 2);
+
+        Py_ssize_t len = PySet_GET_SIZE(o);
+        if (len == 0) {  // empty frozenset should not be re-created.
+            return key;
+        }
+        PyObject *tuple = PyTuple_New(len);
+        if (tuple == NULL) {
+            Py_DECREF(key);
+            return NULL;
+        }
+        Py_ssize_t i = 0, pos = 0;
+        PyObject *item;
+        Py_hash_t hash;
+        while (_PySet_NextEntry(o, &pos, &item, &hash)) {
+            PyObject *k = merge_consts_recursive(c, item);
+            if (k == NULL) {
+                Py_DECREF(tuple);
+                Py_DECREF(key);
+                return NULL;
+            }
+            PyObject *u;
+            if (PyTuple_CheckExact(k)) {
+                u = PyTuple_GET_ITEM(k, 1);
+                Py_INCREF(u);
+                Py_DECREF(k);
+            }
+            else {
+                u = k;
+            }
+            PyTuple_SET_ITEM(tuple, i, u);  // Steals reference of u.
+            i++;
+        }
+
+        // Instead of rewriting o, we create new frozenset and embed in the
+        // key tuple.  Caller should get merged frozenset from the key tuple.
+        PyObject *new = PyFrozenSet_New(tuple);
+        Py_DECREF(tuple);
+        if (new == NULL) {
+            Py_DECREF(key);
+            return NULL;
+        }
+        assert(PyTuple_GET_ITEM(key, 1) == o);
+        Py_DECREF(o);
+        PyTuple_SET_ITEM(key, 1, new);
+    }
 
     return key;
 }
