Skip to content

Commit d753d8a

Browse files
Zheaolitomasr8
andauthored
GH-131798: Narrow the result of _CONTAINS_OP_DICT to bool in the JIT (GH-132269)
Co-authored-by: Tomas R. <[email protected]>
1 parent 71009cb commit d753d8a

File tree

4 files changed

+36
-3
lines changed

4 files changed

+36
-3
lines changed

Lib/test/test_capi/test_opt.py

+27
Original file line numberDiff line numberDiff line change
@@ -1629,6 +1629,33 @@ def testfunc(n):
16291629
self.assertIn("_CONTAINS_OP_SET", uops)
16301630
self.assertNotIn("_TO_BOOL_BOOL", uops)
16311631

1632+
def test_to_bool_bool_contains_op_dict(self):
1633+
"""
1634+
Test that _TO_BOOL_BOOL is removed from code like:
1635+
1636+
res = foo in some_dict
1637+
if res:
1638+
....
1639+
1640+
"""
1641+
def testfunc(n):
1642+
x = 0
1643+
s = {1: 1, 2: 2, 3: 3}
1644+
for _ in range(n):
1645+
a = 2
1646+
in_dict = a in s
1647+
if in_dict:
1648+
x += 1
1649+
return x
1650+
1651+
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
1652+
self.assertEqual(res, TIER2_THRESHOLD)
1653+
self.assertIsNotNone(ex)
1654+
uops = get_opnames(ex)
1655+
self.assertIn("_CONTAINS_OP_DICT", uops)
1656+
self.assertNotIn("_TO_BOOL_BOOL", uops)
1657+
1658+
16321659
def test_remove_guard_for_known_type_str(self):
16331660
def f(n):
16341661
for i in range(n):
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Allow the JIT to remove an extra ``_TO_BOOL_BOOL`` instruction after
2+
``_CONTAINS_OP_DICT`` by setting the return type to bool.

Python/optimizer_bytecodes.c

+4
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,10 @@ dummy_func(void) {
485485
res = sym_new_type(ctx, &PyBool_Type);
486486
}
487487

488+
op(_CONTAINS_OP_DICT, (left, right -- res)) {
489+
res = sym_new_type(ctx, &PyBool_Type);
490+
}
491+
488492
op(_LOAD_CONST, (-- value)) {
489493
PyObject *val = PyTuple_GET_ITEM(co->co_consts, this_instr->oparg);
490494
int opcode = _Py_IsImmortal(val) ? _LOAD_CONST_INLINE_BORROW : _LOAD_CONST_INLINE;

Python/optimizer_cases.c.h

+3-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)