Skip to content

Commit 5847931

Browse files
kumaraditya303kumaraditya303
authored
gh-143732: allow dict subclasses to be specialized (GH-148128)
1 parent de1769f commit 5847931

15 files changed

Lines changed: 455 additions & 200 deletions

‎Include/internal/pycore_dict.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@ extern int _PyDict_Next(
4444

4545
extern int _PyDict_HasOnlyStringKeys(PyObject *mp);
4646

47+
PyAPI_FUNC(PyObject *) _PyDict_Subscript(PyObject *self, PyObject *key);
48+
PyAPI_FUNC(PyObject *) _PyDict_SubscriptKnownHash(PyObject *self, PyObject *key, Py_hash_t hash);
49+
PyAPI_FUNC(int) _PyDict_StoreSubscript(PyObject *self, PyObject *key, PyObject *value);
50+
4751
// Export for '_ctypes' shared extension
4852
PyAPI_FUNC(Py_ssize_t) _PyDict_SizeOf(PyDictObject *);
4953

‎Include/internal/pycore_opcode_metadata.h‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_uop_ids.h‎

Lines changed: 15 additions & 15 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_uop_metadata.h‎

Lines changed: 32 additions & 32 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Lib/test/test_capi/test_opt.py‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2273,10 +2273,49 @@ def f(n):
22732273
self.assertEqual(res, TIER2_THRESHOLD)
22742274
self.assertIsNotNone(ex)
22752275
uops = get_opnames(ex)
2276-
self.assertEqual(uops.count("_GUARD_NOS_DICT"), 0)
2277-
self.assertEqual(uops.count("_STORE_SUBSCR_DICT_KNOWN_HASH"), 1)
2276+
self.assertEqual(uops.count("_GUARD_NOS_DICT_SUBSCRIPT"), 0)
2277+
self.assertEqual(uops.count("_GUARD_NOS_DICT_STORE_SUBSCRIPT"), 0)
22782278
self.assertEqual(uops.count("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH"), 1)
22792279

2280+
def test_dict_subclass_subscr(self):
2281+
import collections
2282+
2283+
def f(n):
2284+
x = 0
2285+
d = collections.defaultdict(int)
2286+
for _ in range(n):
2287+
d["key"] = 1
2288+
x += d["key"]
2289+
return x
2290+
2291+
res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
2292+
self.assertEqual(res, TIER2_THRESHOLD)
2293+
self.assertIsNotNone(ex)
2294+
uops = get_opnames(ex)
2295+
self.assertEqual(uops.count("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH"), 1)
2296+
self.assertEqual(uops.count("_STORE_SUBSCR_DICT_KNOWN_HASH"), 1)
2297+
self.assertEqual(uops.count("_GUARD_NOS_DICT_SUBSCRIPT"), 0)
2298+
self.assertEqual(uops.count("_GUARD_NOS_DICT_STORE_SUBSCRIPT"), 0)
2299+
self.assertEqual(uops.count("_GUARD_TYPE"), 1)
2300+
2301+
def test_dict_subclass_subscr_with_override(self):
2302+
class MyDict(dict):
2303+
def __getitem__(self, key):
2304+
return 42
2305+
2306+
def f(n):
2307+
d = MyDict()
2308+
x = 0
2309+
for _ in range(n):
2310+
x += d["anything"]
2311+
return x
2312+
2313+
res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
2314+
self.assertEqual(res, 42 * TIER2_THRESHOLD)
2315+
self.assertIsNotNone(ex)
2316+
uops = get_opnames(ex)
2317+
self.assertEqual(uops.count("_BINARY_OP_SUBSCR_INIT_CALL"), 1)
2318+
22802319
def test_remove_guard_for_known_type_list(self):
22812320
def f(n):
22822321
x = 0
@@ -2599,6 +2638,23 @@ def testfunc(n):
25992638
self.assertIn("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH", uops)
26002639
self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops)
26012640

2641+
def test_binary_op_subscr_defaultdict_known_hash(self):
2642+
# str, int, bytes, float, complex, tuple and any python object which has generic hash
2643+
import collections
2644+
2645+
def testfunc(n):
2646+
x = 0
2647+
d = collections.defaultdict(lambda: 1)
2648+
for _ in range(n):
2649+
x += d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY] + d[1.5] + d[1+2j]
2650+
return x
2651+
2652+
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
2653+
self.assertEqual(res, 7 * TIER2_THRESHOLD)
2654+
self.assertIsNotNone(ex)
2655+
uops = get_opnames(ex)
2656+
self.assertIn("_BINARY_OP_SUBSCR_DICT_KNOWN_HASH", uops)
2657+
self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops)
26022658

26032659
def test_binary_op_subscr_constant_frozendict_known_hash(self):
26042660
def testfunc(n):
@@ -2635,6 +2691,28 @@ def testfunc(n):
26352691
self.assertIn("_STORE_SUBSCR_DICT_KNOWN_HASH", uops)
26362692
self.assertNotIn("_STORE_SUBSCR_DICT", uops)
26372693

2694+
def test_store_subscr_defaultdict_known_hash(self):
2695+
import collections
2696+
2697+
def testfunc(n):
2698+
d = collections.defaultdict(lambda: 0)
2699+
for _ in range(n):
2700+
d['a'] += 1
2701+
d[1] += 2
2702+
d[b'b'] += 3
2703+
d[(1, 2)] += 4
2704+
d[_GENERIC_KEY] += 5
2705+
d[1.5] += 6
2706+
d[1+2j] += 7
2707+
return d['a'] + d[1] + d[b'b'] + d[(1, 2)] + d[_GENERIC_KEY] + d[1.5] + d[1+2j]
2708+
2709+
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
2710+
self.assertEqual(res, 28 * TIER2_THRESHOLD)
2711+
self.assertIsNotNone(ex)
2712+
uops = get_opnames(ex)
2713+
self.assertIn("_STORE_SUBSCR_DICT_KNOWN_HASH", uops)
2714+
self.assertNotIn("_STORE_SUBSCR_DICT", uops)
2715+
26382716
def test_contains_op(self):
26392717
def testfunc(n):
26402718
x = 0

‎Lib/test/test_opcache.py‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import collections
12
import copy
23
import pickle
34
import dis
@@ -1863,7 +1864,43 @@ class MyFrozenDict(frozendict):
18631864
self.assertEqual(a[2], 3)
18641865

18651866
binary_subscr_frozen_dict_subclass()
1866-
self.assert_no_opcode(binary_subscr_frozen_dict_subclass, "BINARY_OP_SUBSCR_DICT")
1867+
self.assert_specialized(binary_subscr_frozen_dict_subclass, "BINARY_OP_SUBSCR_DICT")
1868+
self.assert_no_opcode(binary_subscr_frozen_dict_subclass, "BINARY_OP")
1869+
1870+
def binary_subscr_defaultdict():
1871+
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
1872+
a = collections.defaultdict(lambda: 42, {1: 2, 2: 3})
1873+
self.assertEqual(a[1], 2)
1874+
self.assertEqual(a[2], 3)
1875+
self.assertEqual(a[7], 42)
1876+
1877+
binary_subscr_defaultdict()
1878+
self.assert_specialized(binary_subscr_defaultdict, "BINARY_OP_SUBSCR_DICT")
1879+
self.assert_no_opcode(binary_subscr_defaultdict, "BINARY_OP")
1880+
1881+
def binary_subscr_counter():
1882+
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
1883+
a = collections.Counter('abcdeabcdabcaba')
1884+
self.assertEqual(a['a'], 5)
1885+
self.assertEqual(a['b'], 4)
1886+
self.assertEqual(a['m'], 0)
1887+
1888+
binary_subscr_counter()
1889+
self.assert_specialized(binary_subscr_counter, "BINARY_OP_SUBSCR_DICT")
1890+
self.assert_no_opcode(binary_subscr_counter, "BINARY_OP")
1891+
1892+
def binary_subscr_dict_subclass_override():
1893+
class MyDict(dict):
1894+
def __getitem__(self, key):
1895+
return 42
1896+
1897+
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
1898+
a = MyDict()
1899+
self.assertEqual(a['a'], 42)
1900+
self.assertEqual(a['b'], 42)
1901+
1902+
binary_subscr_dict_subclass_override()
1903+
self.assert_no_opcode(binary_subscr_dict_subclass_override, "BINARY_OP_SUBSCR_DICT")
18671904

18681905
def binary_subscr_str_int():
18691906
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
@@ -1924,6 +1961,29 @@ def store_subscr_frozen_dict():
19241961
self.assert_specialized(store_subscr_frozen_dict, "STORE_SUBSCR_DICT")
19251962
self.assert_no_opcode(store_subscr_frozen_dict, "STORE_SUBSCR")
19261963

1964+
def store_subscr_defaultdict():
1965+
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
1966+
a = collections.defaultdict(int)
1967+
a[1] = 4
1968+
self.assertEqual(a[1], 4)
1969+
1970+
store_subscr_defaultdict()
1971+
self.assert_specialized(store_subscr_defaultdict, "STORE_SUBSCR_DICT")
1972+
self.assert_no_opcode(store_subscr_defaultdict, "STORE_SUBSCR")
1973+
1974+
def store_subscr_dict_subclass_override():
1975+
class MyDict(dict):
1976+
def __setitem__(self, key, value):
1977+
super().__setitem__(key, value * 2)
1978+
1979+
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
1980+
a = MyDict()
1981+
a['x'] = 5
1982+
self.assertEqual(a['x'], 10)
1983+
1984+
store_subscr_dict_subclass_override()
1985+
self.assert_no_opcode(store_subscr_dict_subclass_override, "STORE_SUBSCR_DICT")
1986+
19271987
@cpython_only
19281988
@requires_specialization
19291989
def test_compare_op(self):

0 commit comments

Comments
 (0)