diff options
| -rw-r--r-- | Lib/test/test_compile.py | 8 | ||||
| -rw-r--r-- | Python/compile.c | 59 |
2 files changed, 65 insertions, 2 deletions
diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index a086ef65b4..56f73f6315 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -615,6 +615,14 @@ if 1: 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 index 7d51819e00..45e78cb22c 100644 --- 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; } |