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

This reverts commit 1005c84.

Author: methane

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}

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;
 }