GH-131798: Type-propagate string/list/tuple slices (GH-134671)

This commit is contained in:
Amit Lavon 2025-06-07 14:08:44 -07:00 committed by GitHub
parent ac9c3431cc
commit 8fdbbf8b18
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 55 additions and 3 deletions

View file

@ -1666,13 +1666,11 @@ class TestUopsOptimization(unittest.TestCase):
self.assertIn("_CONTAINS_OP_DICT", uops)
self.assertNotIn("_TO_BOOL_BOOL", uops)
def test_remove_guard_for_known_type_str(self):
def f(n):
for i in range(n):
false = i == TIER2_THRESHOLD
empty = "X"[:false]
empty += "" # Make JIT realize this is a string.
if empty:
return 1
return 0
@ -2249,6 +2247,34 @@ class TestUopsOptimization(unittest.TestCase):
self.assertNotIn("_LOAD_ATTR_METHOD_NO_DICT", uops)
self.assertNotIn("_LOAD_ATTR_METHOD_LAZY_DICT", uops)
def test_remove_guard_for_slice_list(self):
def f(n):
for i in range(n):
false = i == TIER2_THRESHOLD
sliced = [1, 2, 3][:false]
if sliced:
return 1
return 0
res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
self.assertEqual(res, 0)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_TO_BOOL_LIST", uops)
self.assertNotIn("_GUARD_TOS_LIST", uops)
def test_remove_guard_for_slice_tuple(self):
def f(n):
for i in range(n):
false = i == TIER2_THRESHOLD
a, b = (1, 2, 3)[: false + 2]
_, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_UNPACK_SEQUENCE_TWO_TUPLE", uops)
self.assertNotIn("_GUARD_TOS_TUPLE", uops)
def global_identity(x):
return x

View file

@ -0,0 +1 @@
Make the JIT optimizer understand that slicing a string/list/tuple returns the same type.

View file

@ -1237,6 +1237,20 @@ dummy_func(void) {
sym_set_const(callable, list_append);
}
op(_BINARY_SLICE, (container, start, stop -- res)) {
// Slicing a string/list/tuple always returns the same type.
PyTypeObject *type = sym_get_type(container);
if (type == &PyUnicode_Type ||
type == &PyList_Type ||
type == &PyTuple_Type)
{
res = sym_new_type(ctx, type);
}
else {
res = sym_new_not_null(ctx);
}
}
// END BYTECODES //
}

View file

@ -568,8 +568,19 @@
}
case _BINARY_SLICE: {
JitOptSymbol *container;
JitOptSymbol *res;
res = sym_new_not_null(ctx);
container = stack_pointer[-3];
PyTypeObject *type = sym_get_type(container);
if (type == &PyUnicode_Type ||
type == &PyList_Type ||
type == &PyTuple_Type)
{
res = sym_new_type(ctx, type);
}
else {
res = sym_new_not_null(ctx);
}
stack_pointer[-3] = res;
stack_pointer += -2;
assert(WITHIN_STACK_BOUNDS());