diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index cab7428e107..e652072d9ad 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -299,8 +299,24 @@ jobs: - name: Run flaky MP CPython tests run: | - target/release/rustpython -m test -j 1 ${{ join(matrix.extra_test_args, ' ') }} --slowest --fail-env-changed --timeout 600 -v ${{ env.FLAKY_MP_TESTS }} + for attempt in $(seq 1 5); do + echo "::group::Attempt ${attempt}" + + set +e + target/release/rustpython -m test -j 1 ${{ join(matrix.extra_test_args, ' ') }} --slowest --fail-env-changed --timeout 600 -v ${{ env.FLAKY_MP_TESTS }} + status=$? + set -e + + echo "::endgroup::" + + if [ $status -eq 0 ]; then + exit 0 + fi + done + + exit 1 timeout-minutes: ${{ matrix.timeout }} + shell: bash env: RUSTPYTHON_SKIP_ENV_POLLUTERS: true @@ -363,12 +379,11 @@ jobs: permissions: contents: read checks: write + issues: write pull-requests: write security-events: write # for zizmor steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - with: - persist-credentials: false - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 @@ -415,7 +430,6 @@ jobs: with: level: warning fail_level: error - cleanup: false miri: if: ${{ !contains(github.event.pull_request.labels.*.name, 'skip:ci') }} diff --git a/Cargo.lock b/Cargo.lock index 20bfc4578ff..16941b8265a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3337,6 +3337,7 @@ version = "0.5.0" dependencies = [ "bitflags 2.11.0", "criterion", + "icu_properties", "num_enum", "optional", "rustpython-wtf8", diff --git a/Lib/test/test_base64.py b/Lib/test/test_base64.py index a6739124571..9fe69f4f692 100644 --- a/Lib/test/test_base64.py +++ b/Lib/test/test_base64.py @@ -1,7 +1,7 @@ +import unittest import base64 import binascii import os -import unittest from array import array from test.support import cpython_only from test.support import os_helper @@ -785,6 +785,19 @@ def test_a85decode_errors(self): self.assertRaises(ValueError, base64.a85decode, b'aaaay', foldspaces=True) + self.assertEqual(base64.a85decode(b"a b\nc", ignorechars=b" \n"), + b'\xc9\x89') + with self.assertRaises(ValueError): + base64.a85decode(b"a b\nc", ignorechars=b"") + with self.assertRaises(ValueError): + base64.a85decode(b"a b\nc", ignorechars=b" ") + with self.assertRaises(ValueError): + base64.a85decode(b"a b\nc", ignorechars=b"\n") + with self.assertRaises(TypeError): + base64.a85decode(b"a b\nc", ignorechars=" \n") + with self.assertRaises(TypeError): + base64.a85decode(b"a b\nc", ignorechars=None) + def test_b85decode_errors(self): illegal = list(range(33)) + \ list(b'"\',./:[\\]') + \ diff --git a/Lib/test/test_cmath.py b/Lib/test/test_cmath.py index a96a5780b31..389a3fa0e0a 100644 --- a/Lib/test/test_cmath.py +++ b/Lib/test/test_cmath.py @@ -406,6 +406,8 @@ def polar_with_errno_set(z): _testcapi.set_errno(0) self.check_polar(polar_with_errno_set) + @unittest.skipIf(sys.platform.startswith("sunos"), + "skipping, see gh-138573") def test_phase(self): self.assertAlmostEqual(phase(0), 0.) self.assertAlmostEqual(phase(1.), 0.) diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 9676aded5d1..94a9ae899b0 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -1084,7 +1084,6 @@ def unused_block_while_else(): self.assertEqual('RETURN_VALUE', opcodes[-1].opname) self.assertEqual(None, opcodes[-1].argval) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 3 != 8 def test_false_while_loop(self): def break_in_while(): while False: @@ -1103,7 +1102,6 @@ def continue_in_while(): self.assertEqual('RETURN_VALUE', opcodes[-1].opname) self.assertEqual(None, opcodes[1].argval) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_consts_in_conditionals(self): def and_true(x): return True and x @@ -2578,7 +2576,6 @@ def test_if_else(self): def test_binop(self): self.check_stack_size("x + " * self.N + "x") - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 101 not less than or equal to 6 def test_list(self): self.check_stack_size("[" + "x, " * self.N + "x]") @@ -2586,7 +2583,6 @@ def test_list(self): def test_tuple(self): self.check_stack_size("(" + "x, " * self.N + "x)") - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 101 not less than or equal to 6 def test_set(self): self.check_stack_size("{" + "x, " * self.N + "x}") diff --git a/Lib/test/test_decorators.py b/Lib/test/test_decorators.py index 3a4fc959f6f..78361be9fa1 100644 --- a/Lib/test/test_decorators.py +++ b/Lib/test/test_decorators.py @@ -1,5 +1,5 @@ import unittest -from types import MethodType + def funcattrs(**kwds): def decorate(func): diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py index ab2290c96b8..615212d6024 100644 --- a/Lib/test/test_descr.py +++ b/Lib/test/test_descr.py @@ -1716,6 +1716,28 @@ class SubSpam(spam.spamlist): pass spam_cm.__get__(None, list) self.assertEqual(str(cm.exception), expected_errmsg) + @support.cpython_only + def test_method_get_meth_method_invalid_type(self): + # gh-146615: method_get() for METH_METHOD descriptors used to pass + # Py_TYPE(type)->tp_name as the %V fallback instead of the separate + # %s argument, causing a missing argument for %s and a crash. + # Verify the error message is correct when __get__() is called with a + # non-type as the second argument. + # + # METH_METHOD|METH_FASTCALL|METH_KEYWORDS is the only flag combination + # that enters the affected branch in method_get(). + import io + + obj = io.StringIO() + descr = io.TextIOBase.read + + with self.assertRaises(TypeError) as cm: + descr.__get__(obj, "not_a_type") + self.assertEqual( + str(cm.exception), + "descriptor 'read' needs a type, not 'str', as arg 2", + ) + def test_staticmethods(self): # Testing static methods... class C(object): @@ -4318,6 +4340,7 @@ class C: C.__name__ = Nasty("abc") C.__name__ = "normal" + @unittest.expectedFailureIf(support.is_android, "TODO: RUSTPYTHON; AssertionError: 'C.__rfloordiv__' != 'C.__floordiv__'") def test_subclass_right_op(self): # Testing correct dispatch of subclass overloading __r__... @@ -5170,6 +5193,28 @@ def foo(self): with self.assertRaisesRegex(NotImplementedError, "BAR"): B().foo + @unittest.expectedFailure # TODO: RUSTPYTHON; Wrong error message + def test_staticmethod_new(self): + class MyStaticMethod(staticmethod): + def __init__(self, func): + pass + def func(): pass + sm = MyStaticMethod(func) + self.assertEqual(repr(sm), '') + self.assertIsNone(sm.__func__) + self.assertIsNone(sm.__wrapped__) + + @unittest.expectedFailure # TODO: RUSTPYTHON; Wrong error message + def test_classmethod_new(self): + class MyClassMethod(classmethod): + def __init__(self, func): + pass + def func(): pass + cm = MyClassMethod(func) + self.assertEqual(repr(cm), '') + self.assertIsNone(cm.__func__) + self.assertIsNone(cm.__wrapped__) + class DictProxyTests(unittest.TestCase): def setUp(self): diff --git a/Lib/test/test_fractions.py b/Lib/test/test_fractions.py index 49fc9c2ba23..cf42b86358d 100644 --- a/Lib/test/test_fractions.py +++ b/Lib/test/test_fractions.py @@ -1669,7 +1669,6 @@ def test_float_format_testfile(self): self.assertEqual(float(format(f, fmt2)), float(rhs)) self.assertEqual(float(format(-f, fmt2)), float('-' + rhs)) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: '%' not supported between instances of 'Fraction' and 'complex' def test_complex_handling(self): # See issue gh-102840 for more details. @@ -1697,7 +1696,6 @@ def test_complex_handling(self): message % ("divmod()", "complex", "Fraction"), divmod, b, a) - @unittest.expectedFailure # TODO: RUSTPYTHON; Wrong error message def test_three_argument_pow(self): message = "unsupported operand type(s) for ** or pow(): '%s', '%s', '%s'" self.assertRaisesMessage(TypeError, diff --git a/Lib/test/test_grammar.py b/Lib/test/test_grammar.py index 77bd5a163ce..91eb6cc58f3 100644 --- a/Lib/test/test_grammar.py +++ b/Lib/test/test_grammar.py @@ -304,7 +304,6 @@ def test_var_annot_syntax_errors(self): " nonlocal x\n" " x: int\n") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_var_annot_basic_semantics(self): # execution order with self.assertRaises(ZeroDivisionError): diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 44addde1948..e7c764815d1 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -734,6 +734,60 @@ def keyfunc(obj): keyfunc.skip = 1 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc) + def test_groupby_reentrant_eq_does_not_crash(self): + # regression test for gh-143543 + class Key: + def __init__(self, do_advance): + self.do_advance = do_advance + + def __eq__(self, other): + if self.do_advance: + self.do_advance = False + next(g) + return NotImplemented + return False + + def keys(): + yield Key(True) + yield Key(False) + + g = itertools.groupby([None, None], keys().send) + next(g) + next(g) # must pass with address sanitizer + + @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 4 != 1 + def test_grouper_reentrant_eq_does_not_crash(self): + # regression test for gh-146613 + grouper_iter = None + + class Key: + __hash__ = None + + def __init__(self, do_advance): + self.do_advance = do_advance + + def __eq__(self, other): + nonlocal grouper_iter + if self.do_advance: + self.do_advance = False + if grouper_iter is not None: + try: + next(grouper_iter) + except StopIteration: + pass + return NotImplemented + return True + + def keyfunc(element): + if element == 0: + return Key(do_advance=True) + return Key(do_advance=False) + + g = itertools.groupby(range(4), keyfunc) + key, grouper_iter = next(g) + items = list(grouper_iter) + self.assertEqual(len(items), 1) + def test_filter(self): self.assertEqual(list(filter(isEven, range(6))), [0,2,4]) self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2]) diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py index fea86fe4308..2a3ce809fb0 100644 --- a/Lib/test/test_named_expressions.py +++ b/Lib/test/test_named_expressions.py @@ -4,40 +4,35 @@ class NamedExpressionInvalidTest(unittest.TestCase): - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_01(self): code = """x := 0""" with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_02(self): code = """x = y := 0""" with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_03(self): code = """y := f(x)""" with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_04(self): code = """y0 = y1 := f(x)""" with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_06(self): code = """((a, b) := (1, 2))""" @@ -68,8 +63,7 @@ def test_named_expression_invalid_10(self): with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_11(self): code = """spam(a=1, b := 2)""" @@ -77,8 +71,7 @@ def test_named_expression_invalid_11(self): "positional argument follows keyword argument"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_12(self): code = """spam(a=1, (b := 2))""" @@ -86,8 +79,7 @@ def test_named_expression_invalid_12(self): "positional argument follows keyword argument"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_13(self): code = """spam(a=1, (b := 2))""" @@ -101,8 +93,7 @@ def test_named_expression_invalid_14(self): with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON: wrong error message - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message def test_named_expression_invalid_15(self): code = """(lambda: x := 1)""" @@ -116,8 +107,7 @@ def test_named_expression_invalid_16(self): with self.assertRaisesRegex(SyntaxError, "invalid syntax"): exec(code, {}, {}) - # TODO: RUSTPYTHON - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON def test_named_expression_invalid_17(self): code = "[i := 0, j := 1 for i, j in [(1, 2), (3, 4)]]" @@ -134,8 +124,70 @@ def test_named_expression_invalid_in_class_body(self): "assignment expression within a comprehension cannot be used in a class body"): exec(code, {}, {}) - # TODO: RUSTPYTHON - @unittest.expectedFailure # wrong error message + def test_named_expression_valid_rebinding_iteration_variable(self): + # This test covers that we can reassign variables + # that are not directly assigned in the + # iterable part of a comprehension. + cases = [ + # Regression tests from https://github.com/python/cpython/issues/87447 + ("Complex expression: c", + "{0}(c := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: d", + "{0}(d := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: e", + "{0}(e := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: f", + "{0}(f := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: g", + "{0}(g := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: h", + "{0}(h := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: i", + "{0}(i := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: j", + "{0}(j := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ] + for test_case, code in cases: + for lpar, rpar in [('(', ')'), ('[', ']'), ('{', '}')]: + code = code.format(lpar, rpar) + with self.subTest(case=test_case, lpar=lpar, rpar=rpar): + # Names used in snippets are not defined, + # but we are fine with it: just must not be a SyntaxError. + # Names used in snippets are not defined, + # but we are fine with it: just must not be a SyntaxError. + with self.assertRaises(NameError): + exec(code, {}) # Module scope + with self.assertRaises(NameError): + exec(code, {}, {}) # Class scope + exec(f"lambda: {code}", {}) # Function scope + + def test_named_expression_invalid_rebinding_iteration_variable(self): + # This test covers that we cannot reassign variables + # that are directly assigned in the iterable part of a comprehension. + cases = [ + # Regression tests from https://github.com/python/cpython/issues/87447 + ("Complex expression: a", "a", + "{0}(a := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ("Complex expression: b", "b", + "{0}(b := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"), + ] + for test_case, target, code in cases: + msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'" + for lpar, rpar in [('(', ')'), ('[', ']'), ('{', '}')]: + code = code.format(lpar, rpar) + with self.subTest(case=test_case, lpar=lpar, rpar=rpar): + # Names used in snippets are not defined, + # but we are fine with it: just must not be a SyntaxError. + # Names used in snippets are not defined, + # but we are fine with it: just must not be a SyntaxError. + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope + + @unittest.expectedFailure # TODO: RUSTPYTHON def test_named_expression_invalid_rebinding_list_comprehension_iteration_variable(self): cases = [ ("Local reuse", 'i', "[i := 0 for i in range(5)]"), @@ -151,7 +203,11 @@ def test_named_expression_invalid_rebinding_list_comprehension_iteration_variabl msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'" with self.subTest(case=case): with self.assertRaisesRegex(SyntaxError, msg): - exec(code, {}, {}) + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope def test_named_expression_invalid_rebinding_list_comprehension_inner_loop(self): cases = [ @@ -190,8 +246,7 @@ def test_named_expression_invalid_list_comprehension_iterable_expression(self): with self.assertRaisesRegex(SyntaxError, msg): exec(f"lambda: {code}", {}) # Function scope - # TODO: RUSTPYTHON - @unittest.expectedFailure + @unittest.expectedFailure # TODO: RUSTPYTHON def test_named_expression_invalid_rebinding_set_comprehension_iteration_variable(self): cases = [ ("Local reuse", 'i', "{i := 0 for i in range(5)}"), @@ -202,12 +257,21 @@ def test_named_expression_invalid_rebinding_set_comprehension_iteration_variable ("Unreachable reuse", 'i', "{False or (i:=0) for i in range(5)}"), ("Unreachable nested reuse", 'i', "{(i, j) for i in range(5) for j in range(5) if True or (i:=10)}"), + # Regression tests from https://github.com/python/cpython/issues/87447 + ("Complex expression: a", "a", + "{(a := 1) for a, (*b, c[d+e::f(g)], h.i) in j}"), + ("Complex expression: b", "b", + "{(b := 1) for a, (*b, c[d+e::f(g)], h.i) in j}"), ] for case, target, code in cases: msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'" with self.subTest(case=case): with self.assertRaisesRegex(SyntaxError, msg): - exec(code, {}, {}) + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope def test_named_expression_invalid_rebinding_set_comprehension_inner_loop(self): cases = [ @@ -246,6 +310,84 @@ def test_named_expression_invalid_set_comprehension_iterable_expression(self): with self.assertRaisesRegex(SyntaxError, msg): exec(f"lambda: {code}", {}) # Function scope + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message + def test_named_expression_invalid_rebinding_dict_comprehension_iteration_variable(self): + cases = [ + ("Key reuse", 'i', "{(i := 0): 1 for i in range(5)}"), + ("Value reuse", 'i', "{1: (i := 0) for i in range(5)}"), + ("Both reuse", 'i', "{(i := 0): (i := 0) for i in range(5)}"), + ("Nested reuse", 'j', "{{(j := 0): 1 for i in range(5)} for j in range(5)}"), + ("Reuse inner loop target", 'j', "{(j := 0): 1 for i in range(5) for j in range(5)}"), + ("Unpacking key reuse", 'i', "{(i := 0): 1 for i, j in {(0, 1)}}"), + ("Unpacking value reuse", 'i', "{1: (i := 0) for i, j in {(0, 1)}}"), + ("Reuse in loop condition", 'i', "{i+1: 1 for i in range(5) if (i := 0)}"), + ("Unreachable reuse", 'i', "{(False or (i:=0)): 1 for i in range(5)}"), + ("Unreachable nested reuse", 'i', + "{i: j for i in range(5) for j in range(5) if True or (i:=10)}"), + # Regression tests from https://github.com/python/cpython/issues/87447 + ("Complex expression: a", "a", + "{(a := 1): 1 for a, (*b, c[d+e::f(g)], h.i) in j}"), + ("Complex expression: b", "b", + "{(b := 1): 1 for a, (*b, c[d+e::f(g)], h.i) in j}"), + ] + for case, target, code in cases: + msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'" + with self.subTest(case=case): + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope + + def test_named_expression_invalid_rebinding_dict_comprehension_inner_loop(self): + cases = [ + ("Inner reuse", 'j', "{i: 1 for i in range(5) if (j := 0) for j in range(5)}"), + ("Inner unpacking reuse", 'j', "{i: 1 for i in range(5) if (j := 0) for j, k in {(0, 1)}}"), + ] + for case, target, code in cases: + msg = f"comprehension inner loop cannot rebind assignment expression target '{target}'" + with self.subTest(case=case): + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope + + def test_named_expression_invalid_dict_comprehension_iterable_expression(self): + cases = [ + ("Top level", "{i: 1 for i in (i := range(5))}"), + ("Inside tuple", "{i: 1 for i in (2, 3, i := range(5))}"), + ("Inside list", "{i: 1 for i in [2, 3, i := range(5)]}"), + ("Different name", "{i: 1 for i in (j := range(5))}"), + ("Lambda expression", "{i: 1 for i in (lambda:(j := range(5)))()}"), + ("Inner loop", "{i: 1 for i in range(5) for j in (i := range(5))}"), + ("Nested comprehension", "{i: 1 for i in {j: 2 for j in (k := range(5))}}"), + ("Nested comprehension condition", "{i: 1 for i in {j: 2 for j in range(5) if (j := True)}}"), + ("Nested comprehension body", "{i: 1 for i in {(j := True) for j in range(5)}}"), + ] + msg = "assignment expression cannot be used in a comprehension iterable expression" + for case, code in cases: + with self.subTest(case=case): + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}) # Module scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(code, {}, {}) # Class scope + with self.assertRaisesRegex(SyntaxError, msg): + exec(f"lambda: {code}", {}) # Function scope + + @unittest.expectedFailure # TODO: RUSTPYTHON; wrong error message + def test_named_expression_invalid_mangled_class_variables(self): + code = """class Foo: + def bar(self): + [[(__x:=2) for _ in range(2)] for __x in range(2)] + """ + + with self.assertRaisesRegex(SyntaxError, + "assignment expression cannot rebind comprehension iteration variable '__x'"): + exec(code, {}, {}) + class NamedExpressionAssignmentTest(unittest.TestCase): @@ -299,7 +441,7 @@ def test_named_expression_assignment_09(self): def test_named_expression_assignment_10(self): if (match := 10) == 10: - pass + self.assertEqual(match, 10) else: self.fail("variable was not assigned using named expression") def test_named_expression_assignment_11(self): @@ -341,7 +483,7 @@ def test_named_expression_assignment_14(self): def test_named_expression_assignment_15(self): while a := False: - pass # This will not run + self.fail("While body executed") # This will not run self.assertEqual(a, False) @@ -622,6 +764,18 @@ def test_named_expression_scope_in_genexp(self): for idx, elem in enumerate(genexp): self.assertEqual(elem, b[idx] + a) + def test_named_expression_scope_mangled_names(self): + class Foo: + def f(self_): + global __x1 + __x1 = 0 + [_Foo__x1 := 1 for a in [2]] + self.assertEqual(__x1, 1) + [__x1 := 2 for a in [3]] + self.assertEqual(__x1, 2) + + Foo().f() + self.assertEqual(_Foo__x1, 2) if __name__ == "__main__": unittest.main() diff --git a/Lib/test/test_patma.py b/Lib/test/test_patma.py index 6ca1fa0ba40..40466ec67ba 100644 --- a/Lib/test/test_patma.py +++ b/Lib/test/test_patma.py @@ -3448,7 +3448,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4, 6, 7]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_default_capture(self): def f(command): # 0 match command.split(): # 1 @@ -3463,7 +3462,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4, 6, 7]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_no_default(self): def f(command): # 0 match command.split(): # 1 @@ -3476,7 +3474,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_only_default_wildcard(self): def f(command): # 0 match command.split(): # 1 @@ -3487,7 +3484,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 3]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 3]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_only_default_capture(self): def f(command): # 0 match command.split(): # 1 diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index e20f712a31a..53ff218c4e1 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -785,7 +785,6 @@ def f(a, b, c): c, b, a = a, b, c self.assertNotInBytecode(f, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_mapping(self): for a, b, c in product("_a", "_b", "_c"): pattern = f"{{'a': {a}, 'b': {b}, 'c': {c}}}" @@ -793,7 +792,6 @@ def test_static_swaps_match_mapping(self): code = compile_pattern_with_fast_locals(pattern) self.assertNotInBytecode(code, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_class(self): forms = [ "C({}, {}, {})", @@ -808,7 +806,6 @@ def test_static_swaps_match_class(self): code = compile_pattern_with_fast_locals(pattern) self.assertNotInBytecode(code, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_sequence(self): swaps = {"*_, b, c", "a, *_, c", "a, b, *_"} forms = ["{}, {}, {}", "{}, {}, *{}", "{}, *{}, {}", "*{}, {}, {}"] @@ -863,7 +860,6 @@ def f(): y = x + x self.assertInBytecode(f, 'LOAD_FAST_BORROW_LOAD_FAST_BORROW') - @unittest.expectedFailure # TODO: RUSTPYTHON; RETURN_VALUE def test_load_fast_unknown_simple(self): def f(): if condition(): @@ -906,7 +902,6 @@ def f5(x=0): self.assertInBytecode(f5, 'LOAD_FAST_BORROW') self.assertNotInBytecode(f5, 'LOAD_FAST_CHECK') - @unittest.expectedFailure # TODO: RUSTPYTHON; RETURN_VALUE def test_load_fast_known_because_already_loaded(self): def f(): if condition(): diff --git a/Lib/test/test_pep646_syntax.py b/Lib/test/test_pep646_syntax.py index d9a0aa9a90e..ca8e7d62057 100644 --- a/Lib/test/test_pep646_syntax.py +++ b/Lib/test/test_pep646_syntax.py @@ -305,11 +305,11 @@ {'args': StarredB} >>> def f3(*args: *b, arg1: int): pass - >>> f3.__annotations__ # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE + >>> f3.__annotations__ {'args': StarredB, 'arg1': } >>> def f4(*args: *b, arg1: int = 2): pass - >>> f4.__annotations__ # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE + >>> f4.__annotations__ {'args': StarredB, 'arg1': } >>> def f5(*args: *b = (1,)): pass # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE diff --git a/Lib/test/test_str.py b/Lib/test/test_str.py index 68037923283..15cee0d3a44 100644 --- a/Lib/test/test_str.py +++ b/Lib/test/test_str.py @@ -792,7 +792,6 @@ def test_isdecimal(self): for ch in ['\U0001D7F6', '\U00011066', '\U000104A0']: self.assertTrue(ch.isdecimal(), '{!a} is decimal.'.format(ch)) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: False != True def test_isdigit(self): super().test_isdigit() self.checkequalnofix(True, '\u2460', 'isdigit') diff --git a/Lib/test/test_sys_settrace.py b/Lib/test/test_sys_settrace.py index f9449c7079a..aa2d54ee16e 100644 --- a/Lib/test/test_sys_settrace.py +++ b/Lib/test/test_sys_settrace.py @@ -1957,8 +1957,6 @@ def test_jump_out_of_finally_block(output): finally: output.append(5) - # TODO: RUSTPYTHON - @unittest.expectedFailure @jump_test(1, 5, [], (ValueError, "into an 'except'")) def test_no_jump_into_bare_except_block(output): output.append(1) @@ -1967,8 +1965,6 @@ def test_no_jump_into_bare_except_block(output): except: output.append(5) - # TODO: RUSTPYTHON - @unittest.expectedFailure @jump_test(1, 5, [], (ValueError, "into an 'except'")) def test_no_jump_into_qualified_except_block(output): output.append(1) diff --git a/Lib/test/test_tstring.py b/Lib/test/test_tstring.py index a1f686c8f56..e6984342403 100644 --- a/Lib/test/test_tstring.py +++ b/Lib/test/test_tstring.py @@ -111,7 +111,6 @@ def test_conversions(self): with self.assertRaises(SyntaxError): eval("t'{num!z}'") - @unittest.expectedFailure # TODO: RUSTPYTHON; ? ++++++ def test_debug_specifier(self): # Test debug specifier value = 42 diff --git a/Lib/test/test_typing.py b/Lib/test/test_typing.py index 448d16f1f4a..3ab16478c6b 100644 --- a/Lib/test/test_typing.py +++ b/Lib/test/test_typing.py @@ -4697,7 +4697,6 @@ class D(Generic[T]): pass with self.assertRaises(TypeError): D[()] - @unittest.expectedFailure # TODO: RUSTPYTHON def test_generic_init_subclass_not_called_error(self): notes = ["Note: this exception may have been caused by " r"'GenericTests.test_generic_init_subclass_not_called_error..Base.__init_subclass__' " diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 1a35d2c23d1..0b82f400dfb 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -17,10 +17,11 @@ use crate::{ unparse::UnparseExpr, }; use alloc::borrow::Cow; +use core::mem; use itertools::Itertools; use malachite_bigint::BigInt; use num_complex::Complex; -use num_traits::{Num, ToPrimitive}; +use num_traits::{Num, ToPrimitive, Zero}; use ruff_python_ast as ast; use ruff_text_size::{Ranged, TextRange, TextSize}; use rustpython_compiler_core::{ @@ -124,8 +125,6 @@ enum SuperCallType<'a> { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum BuiltinGeneratorCallKind { Tuple, - List, - Set, All, Any, } @@ -167,6 +166,12 @@ struct Compiler { /// When > 0, the compiler walks AST (consuming sub_tables) but emits no bytecode. /// Mirrors CPython's `c_do_not_emit_bytecode`. do_not_emit_bytecode: u32, + /// Disable constant BoolOp folding in contexts where CPython preserves + /// short-circuit structure, such as starred unpack expressions. + disable_const_boolop_folding: bool, + /// Disable constant tuple/list/set collection folding in contexts where + /// CPython keeps the builder form for later assignment lowering. + disable_const_collection_folding: bool, } #[derive(Clone, Copy)] @@ -439,6 +444,34 @@ enum CollectionType { } impl Compiler { + fn constant_truthiness(constant: &ConstantData) -> bool { + match constant { + ConstantData::Tuple { elements } | ConstantData::Frozenset { elements } => { + !elements.is_empty() + } + ConstantData::Integer { value } => !value.is_zero(), + ConstantData::Float { value } => *value != 0.0, + ConstantData::Complex { value } => value.re != 0.0 || value.im != 0.0, + ConstantData::Boolean { value } => *value, + ConstantData::Str { value } => !value.is_empty(), + ConstantData::Bytes { value } => !value.is_empty(), + ConstantData::Code { .. } | ConstantData::Slice { .. } | ConstantData::Ellipsis => true, + ConstantData::None => false, + } + } + + fn constant_expr_truthiness(&mut self, expr: &ast::Expr) -> CompileResult> { + Ok(self + .try_fold_constant_expr(expr)? + .map(|constant| Self::constant_truthiness(&constant))) + } + + fn disable_load_fast_borrow_for_block(&mut self, block: BlockIdx) { + if block != BlockIdx::NULL { + self.current_code_info().blocks[block.idx()].disable_load_fast_borrow = true; + } + } + fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self { let module_code = ir::CodeInfo { flags: bytecode::CodeFlags::NEWLOCALS, @@ -446,6 +479,7 @@ impl Compiler { private: None, blocks: vec![ir::Block::default()], current_block: BlockIdx::new(0), + annotations_blocks: None, metadata: ir::CodeUnitMetadata { name: code_name.clone(), qualname: Some(code_name), @@ -485,9 +519,62 @@ impl Compiler { in_annotation: false, interactive: false, do_not_emit_bytecode: 0, + disable_const_boolop_folding: false, + disable_const_collection_folding: false, } } + fn compile_expression_without_const_boolop_folding( + &mut self, + expression: &ast::Expr, + ) -> CompileResult<()> { + let previous = self.disable_const_boolop_folding; + self.disable_const_boolop_folding = true; + let result = self.compile_expression(expression); + self.disable_const_boolop_folding = previous; + result.map(|_| ()) + } + + fn compile_expression_without_const_collection_folding( + &mut self, + expression: &ast::Expr, + ) -> CompileResult<()> { + let previous = self.disable_const_collection_folding; + self.disable_const_collection_folding = true; + let result = self.compile_expression(expression); + self.disable_const_collection_folding = previous; + result.map(|_| ()) + } + + fn is_unpack_assignment_target(target: &ast::Expr) -> bool { + matches!(target, ast::Expr::List(_) | ast::Expr::Tuple(_)) + } + + fn compile_module_annotation_setup_sequence( + &mut self, + body: &[ast::Stmt], + ) -> CompileResult<()> { + let (saved_blocks, saved_current_block) = { + let code = self.current_code_info(); + ( + mem::replace(&mut code.blocks, vec![ir::Block::default()]), + mem::replace(&mut code.current_block, BlockIdx::new(0)), + ) + }; + + let result = self.compile_module_annotate(body); + + let annotations_blocks = { + let code = self.current_code_info(); + let annotations_blocks = mem::replace(&mut code.blocks, saved_blocks); + code.current_block = saved_current_block; + annotations_blocks + }; + self.current_code_info().annotations_blocks = Some(annotations_blocks); + + result.map(|_| ()) + } + /// Compile just start and stop of a slice (for BINARY_SLICE/STORE_SLICE) // = codegen_slice_two_parts fn compile_slice_two_parts(&mut self, s: &ast::ExprSlice) -> CompileResult<()> { @@ -584,12 +671,18 @@ impl Compiler { _ => n > 4, }; - // Fold all-constant collections (>= 3 elements) regardless of size - if !seen_star + let can_fold_const_collection = match collection_type { + CollectionType::Tuple => n > 0, + // Match CPython's constant ordering for list/set literals by + // letting the late IR folding passes introduce their tuple-backed + // constants instead of inserting them during AST lowering. + CollectionType::List | CollectionType::Set => false, + }; + if !self.disable_const_collection_folding + && !seen_star && pushed == 0 - && n >= 3 - && elts.iter().all(|e| e.is_constant()) - && let Some(folded) = self.try_fold_constant_collection(elts)? + && can_fold_const_collection + && let Some(folded) = self.try_fold_constant_collection(elts, collection_type)? { match collection_type { CollectionType::Tuple => { @@ -630,9 +723,37 @@ impl Compiler { } // Has stars or too big: use streaming approach + let stream_big_nonconst_collection = if !seen_star + && big + && matches!(collection_type, CollectionType::List | CollectionType::Set) + { + elts.iter().try_fold(false, |has_nonconst, elt| { + if has_nonconst { + return Ok(true); + } + Ok(self.try_fold_constant_expr(elt)?.is_none()) + })? + } else { + false + }; + let mut sequence_built = false; let mut i = 0u32; + if stream_big_nonconst_collection { + match collection_type { + CollectionType::List => { + emit!(self, Instruction::BuildList { count: pushed }); + sequence_built = true; + } + CollectionType::Set => { + emit!(self, Instruction::BuildSet { count: pushed }); + sequence_built = true; + } + CollectionType::Tuple => {} + } + } + for elt in elts.iter() { if let ast::Expr::Starred(ast::ExprStarred { value, .. }) = elt { // When we hit first star, build sequence with elements so far @@ -652,7 +773,7 @@ impl Compiler { } // Compile the starred expression and extend - self.compile_expression(value)?; + self.compile_expression_without_const_boolop_folding(value)?; match collection_type { CollectionType::List => { emit!(self, Instruction::ListExtend { i: 1 }); @@ -1044,17 +1165,18 @@ impl Compiler { /// PEP 709: Inline comprehensions in function-like scopes. /// TODO: Module/class scope inlining needs more work (Cell name resolution edge cases). /// Generator expressions are never inlined. - fn is_inlined_comprehension_context(&self, comprehension_type: ComprehensionType) -> bool { + fn is_inlined_comprehension_context( + &self, + comprehension_type: ComprehensionType, + comp_table: &SymbolTable, + ) -> bool { if comprehension_type == ComprehensionType::Generator { return false; } if !self.ctx.in_func() { return false; } - self.symbol_table_stack - .last() - .and_then(|t| t.sub_tables.get(t.next_sub_table)) - .is_some_and(|st| st.comp_inlined) + comp_table.comp_inlined } /// Enter a new scope @@ -1129,10 +1251,8 @@ impl Compiler { cellvar_cache.insert("__classdict__".to_string()); } - // Handle implicit __conditional_annotations__ cell if needed - if ste.has_conditional_annotations - && matches!(scope_type, CompilerScope::Class | CompilerScope::Module) - { + // Handle implicit __conditional_annotations__ cell if needed. + if Self::scope_needs_conditional_annotations_cell(ste) { cellvar_cache.insert("__conditional_annotations__".to_string()); } @@ -1203,6 +1323,7 @@ impl Compiler { private, blocks: vec![ir::Block::default()], current_block: BlockIdx::new(0), + annotations_blocks: None, metadata: ir::CodeUnitMetadata { name: name.to_owned(), qualname: None, // Will be set below @@ -1238,25 +1359,7 @@ impl Compiler { self.set_qualname(); } - // Emit COPY_FREE_VARS first, then MAKE_CELL (CPython order) - { - let nfrees = self.code_stack.last().unwrap().metadata.freevars.len(); - if nfrees > 0 { - emit!( - self, - Instruction::CopyFreeVars { - n: u32::try_from(nfrees).expect("too many freevars"), - } - ); - } - } - { - let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); - for i in 0..ncells { - let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); - emit!(self, Instruction::MakeCell { i: i_varnum }); - } - } + self.emit_prefix_cell_setup(); // Emit RESUME (handles async preamble and module lineno 0) // CPython: LOCATION(lineno, lineno, 0, 0), then loc.lineno = 0 for module @@ -1302,11 +1405,42 @@ impl Compiler { location, end_location, except_handler, + folded_from_nonliteral_expr: false, lineno_override, cache_entries: 0, }); } + fn emit_prefix_cell_setup(&mut self) { + let metadata = &self.code_stack.last().unwrap().metadata; + let varnames = metadata.varnames.clone(); + let cellvars = metadata.cellvars.clone(); + let freevars = metadata.freevars.clone(); + let ncells = cellvars.len(); + if ncells > 0 { + let cellfixedoffsets = ir::build_cellfixedoffsets(&varnames, &cellvars, &freevars); + let mut sorted = vec![None; varnames.len() + ncells]; + for (oldindex, fixed) in cellfixedoffsets.iter().copied().take(ncells).enumerate() { + sorted[fixed as usize] = Some(oldindex); + } + for oldindex in sorted.into_iter().flatten() { + let i_varnum: oparg::VarNum = + u32::try_from(oldindex).expect("too many cellvars").into(); + emit!(self, Instruction::MakeCell { i: i_varnum }); + } + } + + let nfrees = freevars.len(); + if nfrees > 0 { + emit!( + self, + Instruction::CopyFreeVars { + n: u32::try_from(nfrees).expect("too many freevars"), + } + ); + } + } + fn push_output( &mut self, flags: bytecode::CodeFlags, @@ -1818,7 +1952,7 @@ impl Compiler { self.future_annotations = symbol_table.future_annotations; // Module-level __conditional_annotations__ cell - let has_module_cond_ann = symbol_table.has_conditional_annotations; + let has_module_cond_ann = Self::scope_needs_conditional_annotations_cell(&symbol_table); if has_module_cond_ann { self.current_code_info() .metadata @@ -1828,16 +1962,14 @@ impl Compiler { self.symbol_table_stack.push(symbol_table); - // Emit MAKE_CELL for module-level cells (before RESUME) + // Match flowgraph.c insert_prefix_instructions() for module-level + // synthetic cells before RESUME. if has_module_cond_ann { - let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); - for i in 0..ncells { - let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); - emit!(self, Instruction::MakeCell { i: i_varnum }); - } + self.emit_prefix_cell_setup(); } self.emit_resume_for_scope(CompilerScope::Module, 1); + emit!(self, PseudoInstruction::AnnotationsPlaceholder); let (doc, statements) = split_doc(&body.body, &self.opts); if let Some(value) = doc { @@ -1848,26 +1980,26 @@ impl Compiler { emit!(self, Instruction::StoreName { namei: doc }) } - // Handle annotations based on future_annotations flag + // Handle annotation bookkeeping in CPython order: initialize the + // conditional annotation set first, then materialize __annotations__. if Self::find_ann(statements) { + if Self::scope_needs_conditional_annotations_cell(self.current_symbol_table()) { + emit!(self, Instruction::BuildSet { count: 0 }); + self.store_name("__conditional_annotations__")?; + } + if self.future_annotations { - // PEP 563: Initialize __annotations__ dict emit!(self, Instruction::SetupAnnotations); - } else { - // PEP 649: Generate __annotate__ function FIRST (before statements) - self.compile_module_annotate(statements)?; - - // PEP 649: Initialize __conditional_annotations__ set after __annotate__ - if self.current_symbol_table().has_conditional_annotations { - emit!(self, Instruction::BuildSet { count: 0 }); - self.store_name("__conditional_annotations__")?; - } } } // Compile all statements self.compile_statements(statements)?; + if Self::find_ann(statements) && !self.future_annotations { + self.compile_module_annotation_setup_sequence(statements)?; + } + assert_eq!(self.code_stack.len(), size_before); // Emit None at end: @@ -1886,6 +2018,7 @@ impl Compiler { self.symbol_table_stack.push(symbol_table); self.emit_resume_for_scope(CompilerScope::Module, 1); + emit!(self, PseudoInstruction::AnnotationsPlaceholder); // Handle annotations based on future_annotations flag if Self::find_ann(body) { @@ -1893,10 +2026,8 @@ impl Compiler { // PEP 563: Initialize __annotations__ dict emit!(self, Instruction::SetupAnnotations); } else { - // PEP 649: Generate __annotate__ function FIRST (before statements) - self.compile_module_annotate(body)?; - - // PEP 649: Initialize __conditional_annotations__ set after __annotate__ + // PEP 649: Initialize __conditional_annotations__ before the body. + // CPython generates __annotate__ after the body in codegen_body(). if self.current_symbol_table().has_conditional_annotations { emit!(self, Instruction::BuildSet { count: 0 }); self.store_name("__conditional_annotations__")?; @@ -1940,6 +2071,10 @@ impl Compiler { self.emit_load_const(ConstantData::None); }; + if Self::find_ann(body) && !self.future_annotations { + self.compile_module_annotation_setup_sequence(body)?; + } + self.emit_return_value(); Ok(()) } @@ -1994,6 +2129,20 @@ impl Compiler { Ok(()) } + fn scope_needs_conditional_annotations_cell(symbol_table: &SymbolTable) -> bool { + match symbol_table.typ { + CompilerScope::Module => { + symbol_table.has_conditional_annotations + || (symbol_table.future_annotations && symbol_table.annotation_block.is_some()) + } + CompilerScope::Class => { + symbol_table.has_conditional_annotations + || symbol_table.lookup("__conditional_annotations__").is_some() + } + _ => false, + } + } + fn load_name(&mut self, name: &str) -> CompileResult<()> { self.compile_name(name, NameUsage::Load) } @@ -2379,7 +2528,7 @@ impl Compiler { let dominated_by_interactive = self.interactive && !self.ctx.in_func() && !self.ctx.in_class; if !dominated_by_interactive && Self::is_const_expression(value) { - // Skip compilation entirely - the expression has no side effects + emit!(self, Instruction::Nop); } else { self.compile_expression(value)?; @@ -2405,8 +2554,9 @@ impl Compiler { .. }) => { self.enter_conditional_block(); - self.compile_if(test, body, elif_else_clauses)?; + self.compile_if(test, body, elif_else_clauses, test.range())?; self.leave_conditional_block(); + self.set_source_range(statement.range()); } ast::Stmt::While(ast::StmtWhile { test, body, orelse, .. @@ -2508,7 +2658,6 @@ impl Compiler { if self.opts.optimize == 0 { let after_block = self.new_block(); self.compile_jump_if(test, true, after_block)?; - emit!( self, Instruction::LoadCommonConstant { @@ -2516,9 +2665,8 @@ impl Compiler { } ); if let Some(e) = msg { - emit!(self, Instruction::PushNull); self.compile_expression(e)?; - emit!(self, Instruction::Call { argc: 1 }); + emit!(self, Instruction::Call { argc: 0 }); } emit!( self, @@ -2526,7 +2674,6 @@ impl Compiler { argc: bytecode::RaiseKind::Raise, } ); - self.switch_to_block(after_block); } else { // Optimized-out asserts still need to consume any nested @@ -2593,7 +2740,11 @@ impl Compiler { self.switch_to_block(dead); } ast::Stmt::Assign(ast::StmtAssign { targets, value, .. }) => { - self.compile_expression(value)?; + if targets.len() == 1 && Self::is_unpack_assignment_target(&targets[0]) { + self.compile_expression_without_const_collection_folding(value)?; + } else { + self.compile_expression(value)?; + } for (i, target) in targets.iter().enumerate() { if i + 1 != targets.len() { @@ -3437,16 +3588,26 @@ impl Compiler { handlers: &[ast::ExceptHandler], orelse: &[ast::Stmt], ) -> CompileResult<()> { + let normal_exit_range = orelse + .last() + .map(ast::Stmt::range) + .or_else(|| body.last().map(ast::Stmt::range)); let handler_block = self.new_block(); let cleanup_block = self.new_block(); let end_block = self.new_block(); - let orelse_block = if orelse.is_empty() { - end_block - } else { - self.new_block() - }; + let has_bare_except = handlers.iter().any(|handler| { + matches!( + handler, + ast::ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { + type_: None, + .. + }) + ) + }); + if has_bare_except { + self.disable_load_fast_borrow_for_block(end_block); + } - emit!(self, Instruction::Nop); emit!( self, PseudoInstruction::SetupFinally { @@ -3459,11 +3620,10 @@ impl Compiler { self.pop_fblock(FBlockType::TryExcept); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); + self.compile_statements(orelse)?; emit!( self, - PseudoInstruction::JumpNoInterrupt { - delta: orelse_block - } + PseudoInstruction::JumpNoInterrupt { delta: end_block } ); self.set_no_location(); @@ -3505,7 +3665,6 @@ impl Compiler { self.store_name(alias.as_str())?; let cleanup_end = self.new_block(); - let handler_normal_exit = self.new_block(); emit!(self, PseudoInstruction::SetupCleanup { delta: cleanup_end }); self.push_fblock_full( FBlockType::HandlerCleanup, @@ -3519,25 +3678,6 @@ impl Compiler { self.pop_fblock(FBlockType::HandlerCleanup); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); - emit!( - self, - PseudoInstruction::JumpNoInterrupt { - delta: handler_normal_exit - } - ); - self.set_no_location(); - - self.switch_to_block(cleanup_end); - self.emit_load_const(ConstantData::None); - self.set_no_location(); - self.store_name(alias.as_str())?; - self.set_no_location(); - self.compile_name(alias.as_str(), NameUsage::Delete)?; - self.set_no_location(); - emit!(self, Instruction::Reraise { depth: 1 }); - self.set_no_location(); - - self.switch_to_block(handler_normal_exit); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); self.pop_fblock(FBlockType::ExceptionHandler); @@ -3556,6 +3696,16 @@ impl Compiler { PseudoInstruction::JumpNoInterrupt { delta: end_block } ); self.set_no_location(); + + self.switch_to_block(cleanup_end); + self.emit_load_const(ConstantData::None); + self.set_no_location(); + self.store_name(alias.as_str())?; + self.set_no_location(); + self.compile_name(alias.as_str(), NameUsage::Delete)?; + self.set_no_location(); + emit!(self, Instruction::Reraise { depth: 1 }); + self.set_no_location(); } else { emit!(self, Instruction::PopTop); self.push_fblock(FBlockType::HandlerCleanup, end_block, end_block)?; @@ -3591,18 +3741,10 @@ impl Compiler { emit!(self, Instruction::Reraise { depth: 1 }); self.set_no_location(); - if !orelse.is_empty() { - self.switch_to_block(orelse_block); - self.set_no_location(); - self.compile_statements(orelse)?; - emit!( - self, - PseudoInstruction::JumpNoInterrupt { delta: end_block } - ); - self.set_no_location(); - } - self.switch_to_block(end_block); + if let Some(range) = normal_exit_range { + self.set_source_range(range); + } Ok(()) } @@ -4224,18 +4366,33 @@ impl Compiler { parameters: &ast::Parameters, returns: Option<&ast::Expr>, ) -> CompileResult { + let has_signature_annotations = parameters + .args + .iter() + .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) + .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) + .chain(parameters.kwarg.as_deref()) + .any(|param| param.annotation.is_some()) + || returns.is_some(); + if !has_signature_annotations { + return Ok(false); + } + // Try to enter annotation scope - returns None if no annotation_block exists let Some(saved_ctx) = self.enter_annotation_scope(func_name)? else { return Ok(false); }; // Count annotations - let parameters_iter = core::iter::empty() - .chain(¶meters.posonlyargs) - .chain(¶meters.args) - .chain(¶meters.kwonlyargs) + let parameters_iter = parameters + .args + .iter() .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.kwarg.as_deref()); let num_annotations: u32 = @@ -4244,12 +4401,13 @@ impl Compiler { + if returns.is_some() { 1 } else { 0 }; // Compile annotations inside the annotation scope - let parameters_iter = core::iter::empty() - .chain(¶meters.posonlyargs) - .chain(¶meters.args) - .chain(¶meters.kwonlyargs) + let parameters_iter = parameters + .args + .iter() .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.kwarg.as_deref()); for param in parameters_iter { @@ -4287,24 +4445,15 @@ impl Compiler { Ok(true) } - /// Collect simple annotations from module body in AST order (including nested blocks) - /// Returns list of (name, annotation_expr) pairs - /// This must match the order that annotations are compiled to ensure - /// conditional_annotation_index stays in sync with __annotate__ enumeration. - fn collect_simple_annotations(body: &[ast::Stmt]) -> Vec<(&str, &ast::Expr)> { - fn walk<'a>(stmts: &'a [ast::Stmt], out: &mut Vec<(&'a str, &'a ast::Expr)>) { + /// Collect annotated assignments from module/class body in AST order + /// (including nested conditional blocks). This preserves the same walk + /// order as symbol-table construction so the annotation scope's + /// `sub_tables` cursor stays aligned. + fn collect_annotations(body: &[ast::Stmt]) -> Vec<&ast::StmtAnnAssign> { + fn walk<'a>(stmts: &'a [ast::Stmt], out: &mut Vec<&'a ast::StmtAnnAssign>) { for stmt in stmts { match stmt { - ast::Stmt::AnnAssign(ast::StmtAnnAssign { - target, - annotation, - simple, - .. - }) if *simple && matches!(target.as_ref(), ast::Expr::Name(_)) => { - if let ast::Expr::Name(ast::ExprName { id, .. }) = target.as_ref() { - out.push((id.as_str(), annotation.as_ref())); - } - } + ast::Stmt::AnnAssign(stmt) => out.push(stmt), ast::Stmt::If(ast::StmtIf { body, elif_else_clauses, @@ -4355,10 +4504,13 @@ impl Compiler { /// Compile module-level __annotate__ function (PEP 649) /// Returns true if __annotate__ was created and stored fn compile_module_annotate(&mut self, body: &[ast::Stmt]) -> CompileResult { - // Collect simple annotations from module body first - let annotations = Self::collect_simple_annotations(body); + let annotations = Self::collect_annotations(body); + let simple_annotation_count = annotations + .iter() + .filter(|stmt| stmt.simple && matches!(stmt.target.as_ref(), ast::Expr::Name(_))) + .count(); - if annotations.is_empty() { + if simple_annotation_count == 0 { return Ok(false); } @@ -4400,20 +4552,42 @@ impl Compiler { // Emit format validation: if format > VALUE_WITH_FAKE_GLOBALS: raise NotImplementedError self.emit_format_validation()?; - if has_conditional { - // PEP 649: Build dict incrementally, checking conditional annotations - // Start with empty dict - emit!(self, Instruction::BuildMap { count: 0 }); + emit!(self, Instruction::BuildMap { count: 0 }); + + let mut simple_idx = 0usize; + for stmt in annotations { + let ast::StmtAnnAssign { + target, + annotation, + simple, + .. + } = stmt; + let simple_name = if *simple { + match target.as_ref() { + ast::Expr::Name(ast::ExprName { id, .. }) => Some(id.as_str()), + _ => None, + } + } else { + None + }; + + if simple_name.is_none() { + if !self.future_annotations { + self.do_not_emit_bytecode += 1; + let result = self.compile_annotation(annotation); + self.do_not_emit_bytecode -= 1; + result?; + } + continue; + } - // Process each annotation - for (idx, (name, annotation)) in annotations.iter().enumerate() { - // Check if index is in __conditional_annotations__ - let not_set_block = self.new_block(); + let not_set_block = has_conditional.then(|| self.new_block()); + let name = simple_name.expect("missing simple annotation name"); - // LOAD_CONST index - self.emit_load_const(ConstantData::Integer { value: idx.into() }); - // Load __conditional_annotations__ from appropriate scope - // Class scope: LoadDeref (freevars), Module scope: LoadGlobal + if has_conditional { + self.emit_load_const(ConstantData::Integer { + value: simple_idx.into(), + }); if parent_scope_type == CompilerScope::Class { let idx = self.get_free_var_index("__conditional_annotations__")?; emit!(self, Instruction::LoadDeref { i: idx }); @@ -4421,60 +4595,35 @@ impl Compiler { let cond_annotations_name = self.name("__conditional_annotations__"); self.emit_load_global(cond_annotations_name, false); } - // CONTAINS_OP (in) emit!( self, Instruction::ContainsOp { invert: bytecode::Invert::No } ); - // POP_JUMP_IF_FALSE not_set emit!( self, Instruction::PopJumpIfFalse { - delta: not_set_block + delta: not_set_block.expect("missing not_set block") } ); - - // Annotation value - self.compile_annotation(annotation)?; - // COPY dict to TOS - emit!(self, Instruction::Copy { i: 2 }); - // LOAD_CONST name - self.emit_load_const(ConstantData::Str { - value: self.mangle(name).into_owned().into(), - }); - // STORE_SUBSCR - dict[name] = value - emit!(self, Instruction::StoreSubscr); - - // not_set label - self.switch_to_block(not_set_block); } - // Return the dict - emit!(self, Instruction::ReturnValue); - } else { - // No conditional annotations - use simple BuildMap - let num_annotations = u32::try_from(annotations.len()).expect("too many annotations"); + self.compile_annotation(annotation)?; + emit!(self, Instruction::Copy { i: 2 }); + self.emit_load_const(ConstantData::Str { + value: self.mangle(name).into_owned().into(), + }); + emit!(self, Instruction::StoreSubscr); + simple_idx += 1; - // Compile annotations inside the annotation scope - for (name, annotation) in annotations { - self.emit_load_const(ConstantData::Str { - value: self.mangle(name).into_owned().into(), - }); - self.compile_annotation(annotation)?; + if let Some(not_set_block) = not_set_block { + self.switch_to_block(not_set_block); } - - // Build the map and return it - emit!( - self, - Instruction::BuildMap { - count: num_annotations, - } - ); - emit!(self, Instruction::ReturnValue); } + emit!(self, Instruction::ReturnValue); + // Exit annotation scope - pop symbol table, restore to parent's annotation_block, and get code let annotation_table = self.pop_symbol_table(); // Restore annotation_block to module's symbol table @@ -5081,55 +5230,37 @@ impl Compiler { } ); - // PEP 649: Initialize __classdict__ cell (before __doc__) + // Set __type_params__ from the enclosing type-params closure when + // compiling a generic class body. + if type_params.is_some() { + self.load_name(".type_params")?; + self.store_name("__type_params__")?; + } + + // PEP 649: Initialize __classdict__ after synthetic generic-class + // setup so nested generic classes match CPython's prologue order. if self.current_symbol_table().needs_classdict { emit!(self, Instruction::LoadLocals); let classdict_idx = self.get_cell_var_index("__classdict__")?; emit!(self, Instruction::StoreDeref { i: classdict_idx }); } - // Store __doc__ only if there's an explicit docstring + // Store __doc__ only if there's an explicit docstring. if let Some(doc) = doc_str { self.emit_load_const(ConstantData::Str { value: doc.into() }); let doc_name = self.name("__doc__"); emit!(self, Instruction::StoreName { namei: doc_name }); } - // Set __type_params__ if we have type parameters - if type_params.is_some() { - // Load .type_params from enclosing scope - let dot_type_params = self.name(".type_params"); - emit!( - self, - Instruction::LoadName { - namei: dot_type_params - } - ); - - // Store as __type_params__ - let dunder_type_params = self.name("__type_params__"); - emit!( - self, - Instruction::StoreName { - namei: dunder_type_params - } - ); - } - - // Handle class annotations based on future_annotations flag + // Handle class annotation bookkeeping in CPython order. if Self::find_ann(body) { + if Self::scope_needs_conditional_annotations_cell(self.current_symbol_table()) { + emit!(self, Instruction::BuildSet { count: 0 }); + self.store_name("__conditional_annotations__")?; + } + if self.future_annotations { - // PEP 563: Initialize __annotations__ dict for class emit!(self, Instruction::SetupAnnotations); - } else { - // PEP 649: Initialize __conditional_annotations__ set if needed for class - if self.current_symbol_table().has_conditional_annotations { - emit!(self, Instruction::BuildSet { count: 0 }); - self.store_name("__conditional_annotations__")?; - } - - // PEP 649: Generate __annotate__ function for class annotations - self.compile_module_annotate(body)?; } } @@ -5146,6 +5277,10 @@ impl Compiler { // 3. Compile the class body self.compile_statements(body)?; + if Self::find_ann(body) && !self.future_annotations { + self.compile_module_annotate(body)?; + } + // 4. Handle __classcell__ if needed let classcell_idx = self .code_stack @@ -5254,15 +5389,10 @@ impl Compiler { in_async_scope: false, }; - // Compile type parameters and store as .type_params + // Compile type parameters and store them in the synthetic cell that + // generic class bodies close over. self.compile_type_params(type_params.unwrap())?; - let dot_type_params = self.name(".type_params"); - emit!( - self, - Instruction::StoreName { - namei: dot_type_params - } - ); + self.store_name(".type_params")?; } // Step 2: Compile class body (always done, whether generic or not) @@ -5278,47 +5408,25 @@ impl Compiler { // Step 3: Generate the rest of the code for the call if is_generic { - // Still in type params scope - let dot_type_params = self.name(".type_params"); - let dot_generic_base = self.name(".generic_base"); - - // Create .generic_base - emit!( - self, - Instruction::LoadName { - namei: dot_type_params - } - ); - emit!( - self, - Instruction::CallIntrinsic1 { - func: bytecode::IntrinsicFunction1::SubscriptGeneric - } - ); - emit!( - self, - Instruction::StoreName { - namei: dot_generic_base - } - ); - // Generate class creation code emit!(self, Instruction::LoadBuildClass); emit!(self, Instruction::PushNull); - // Set up the class function with type params - let mut func_flags = bytecode::MakeFunctionFlags::new(); + // Create the class body function with the .type_params closure + // captured through the class code object's freevars. + self.make_closure(class_code, bytecode::MakeFunctionFlags::new())?; + self.emit_load_const(ConstantData::Str { value: name.into() }); + + // Create .generic_base after the class function and name are on the + // stack so the remaining call shape matches CPython's ordering. + self.load_name(".type_params")?; emit!( self, - Instruction::LoadName { - namei: dot_type_params + Instruction::CallIntrinsic1 { + func: bytecode::IntrinsicFunction1::SubscriptGeneric } ); - func_flags.insert(bytecode::MakeFunctionFlag::TypeParams); - - // Create class function with closure - self.make_closure(class_code, func_flags)?; - self.emit_load_const(ConstantData::Str { value: name.into() }); + self.store_name(".generic_base")?; // Compile bases and call __build_class__ // Check for starred bases or **kwargs @@ -5354,12 +5462,7 @@ impl Compiler { } // Add .generic_base as final element - emit!( - self, - Instruction::LoadName { - namei: dot_generic_base - } - ); + self.load_name(".generic_base")?; emit!(self, Instruction::ListAppend { i: 1 }); // Convert list to tuple @@ -5390,12 +5493,7 @@ impl Compiler { }; // Load .generic_base as the last base - emit!( - self, - Instruction::LoadName { - namei: dot_generic_base - } - ); + self.load_name(".generic_base")?; let nargs = 2 + u32::try_from(base_count).expect("too many base classes") + 1; @@ -5468,84 +5566,36 @@ impl Compiler { test: &ast::Expr, body: &[ast::Stmt], elif_else_clauses: &[ast::ElifElseClause], + _stmt_range: TextRange, ) -> CompileResult<()> { - let constant = Self::expr_constant(test); + let end_block = self.new_block(); + let next_block = if elif_else_clauses.is_empty() { + end_block + } else { + self.new_block() + }; - // If the test is constant false, walk the body (consuming sub_tables) - // but don't emit bytecode - if constant == Some(false) { - self.emit_nop(); - self.do_not_emit_bytecode += 1; - self.compile_statements(body)?; - self.do_not_emit_bytecode -= 1; - // Compile the elif/else chain (if any) - match elif_else_clauses { - [] => {} - [first, rest @ ..] => { - if let Some(elif_test) = &first.test { - self.compile_if(elif_test, &first.body, rest)?; - } else { - self.compile_statements(&first.body)?; - } - } - } - return Ok(()); + if matches!(self.constant_expr_truthiness(test)?, Some(false)) { + self.disable_load_fast_borrow_for_block(next_block); } + self.compile_jump_if(test, false, next_block)?; + self.compile_statements(body)?; - // If the test is constant true, compile body directly, - // but walk elif/else without emitting (including elif tests to consume sub_tables) - if constant == Some(true) { - self.emit_nop(); - self.compile_statements(body)?; - self.do_not_emit_bytecode += 1; - for clause in elif_else_clauses { - if let Some(elif_test) = &clause.test { - self.compile_expression(elif_test)?; - } - self.compile_statements(&clause.body)?; - } - self.do_not_emit_bytecode -= 1; + let Some((clause, rest)) = elif_else_clauses.split_first() else { + self.switch_to_block(end_block); return Ok(()); - } - - // Non-constant test: normal compilation - match elif_else_clauses { - // Only if - [] => { - let after_block = self.new_block(); - self.compile_jump_if(test, false, after_block)?; - self.compile_statements(body)?; - self.switch_to_block(after_block); - } - // If, elif*, elif/else - [rest @ .., tail] => { - let after_block = self.new_block(); - let mut next_block = self.new_block(); - - self.compile_jump_if(test, false, next_block)?; - self.compile_statements(body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); + }; - for clause in rest { - self.switch_to_block(next_block); - next_block = self.new_block(); - if let Some(test) = &clause.test { - self.compile_jump_if(test, false, next_block)?; - } else { - unreachable!() // must be elif - } - self.compile_statements(&clause.body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); - } + emit!(self, PseudoInstruction::Jump { delta: end_block }); + self.switch_to_block(next_block); - self.switch_to_block(next_block); - if let Some(test) = &tail.test { - self.compile_jump_if(test, false, after_block)?; - } - self.compile_statements(&tail.body)?; - self.switch_to_block(after_block); - } + if let Some(test) = &clause.test { + self.compile_if(test, &clause.body, rest, test.range())?; + } else { + debug_assert!(rest.is_empty()); + self.compile_statements(&clause.body)?; } + self.switch_to_block(end_block); Ok(()) } @@ -5557,37 +5607,17 @@ impl Compiler { ) -> CompileResult<()> { self.enter_conditional_block(); - let constant = Self::expr_constant(test); - - // while False: body → walk body (consuming sub_tables) but don't emit, - // then compile orelse - if constant == Some(false) { - self.emit_nop(); - let while_block = self.new_block(); - let after_block = self.new_block(); - self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; - self.do_not_emit_bytecode += 1; - self.compile_statements(body)?; - self.do_not_emit_bytecode -= 1; - self.pop_fblock(FBlockType::WhileLoop); - self.compile_statements(orelse)?; - self.leave_conditional_block(); - return Ok(()); - } - let while_block = self.new_block(); let else_block = self.new_block(); let after_block = self.new_block(); self.switch_to_block(while_block); self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; - - // while True: → no condition test, just NOP - if constant == Some(true) { - self.emit_nop(); - } else { - self.compile_jump_if(test, false, else_block)?; + if matches!(self.constant_expr_truthiness(test)?, Some(false)) { + self.disable_load_fast_borrow_for_block(else_block); + self.disable_load_fast_borrow_for_block(after_block); } + self.compile_jump_if(test, false, else_block)?; let was_in_loop = self.ctx.loop_data.replace((while_block, after_block)); self.compile_statements(body)?; @@ -5855,24 +5885,7 @@ impl Compiler { let mut end_async_for_target = BlockIdx::NULL; // The thing iterated: - // Optimize: `for x in [a, b, c]` → use tuple instead of list - // Skip for async-for (GET_AITER expects the original type) - if !is_async - && let ast::Expr::List(ast::ExprList { elts, .. }) = iter - && !elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))) - { - for elt in elts { - self.compile_expression(elt)?; - } - emit!( - self, - Instruction::BuildTuple { - count: u32::try_from(elts.len()).expect("too many elements"), - } - ); - } else { - self.compile_expression(iter)?; - } + self.compile_for_iterable_expression(iter, is_async)?; if is_async { if self.ctx.func != FunctionContext::AsyncFunction { @@ -5907,8 +5920,13 @@ impl Compiler { emit!(self, Instruction::ForIter { delta: else_block }); - // Start of loop iteration, set targets: + // Match CPython's line attribution by compiling the loop target on + // the target range directly instead of leaving a synthetic anchor + // NOP between FOR_ITER and the unpack/store sequence. + let saved_range = self.current_source_range; + self.set_source_range(target.range()); self.compile_store(target)?; + self.set_source_range(saved_range); }; let was_in_loop = self.ctx.loop_data.replace((for_block, after_block)); @@ -5943,6 +5961,38 @@ impl Compiler { Ok(()) } + fn compile_for_iterable_expression( + &mut self, + iter: &ast::Expr, + is_async: bool, + ) -> CompileResult<()> { + // Match CPython's iterable lowering for `for`/comprehension fronts: + // a non-starred list literal used only for iteration is emitted as a tuple. + // Skip async-for/async comprehension iteration because GET_AITER expects + // the original object semantics. + if !is_async + && let ast::Expr::List(ast::ExprList { elts, .. }) = iter + && !elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))) + { + if let Some(folded) = self.try_fold_constant_collection(elts, CollectionType::List)? { + self.emit_load_const(folded); + } else { + for elt in elts { + self.compile_expression(elt)?; + } + emit!( + self, + Instruction::BuildTuple { + count: u32::try_from(elts.len()).expect("too many elements"), + } + ); + } + return Ok(()); + } + + self.compile_expression(iter) + } + fn forbidden_name(&mut self, name: &str, ctx: NameUsage) -> CompileResult { if ctx == NameUsage::Store && name == "__debug__" { return Err(self.error(CodegenErrorType::Assign("__debug__"))); @@ -6150,9 +6200,17 @@ impl Compiler { // Keep the subject around for extracting elements. pc.on_top += 1; for (i, pattern) in patterns.iter().enumerate() { - // if pattern.is_wildcard() { - // continue; - // } + let is_true_wildcard = matches!( + pattern, + ast::Pattern::MatchAs(ast::PatternMatchAs { + pattern: None, + name: None, + .. + }) + ); + if is_true_wildcard { + continue; + } if i == star { // This must be a starred wildcard. // assert!(pattern.is_star_wildcard()); @@ -6660,6 +6718,7 @@ impl Compiler { pc.stores.insert(insert_pos + j, elem); } // Also perform the same rotation on the evaluation stack. + self.set_source_range(alt.range()); for _ in 0..=i_stores { self.pattern_helper_rotate(i_control + 1)?; } @@ -6668,7 +6727,9 @@ impl Compiler { } } // Emit a jump to the common end label and reset any failure jump targets. + self.set_source_range(alt.range()); emit!(self, PseudoInstruction::Jump { delta: end }); + self.set_source_range(alt.range()); self.emit_and_reset_fail_pop(pc)?; } @@ -6680,6 +6741,7 @@ impl Compiler { // In Rust, old_pc is a local clone, so we need not worry about that. // No alternative matched: pop the subject and fail. + self.set_source_range(p.range()); emit!(self, Instruction::PopTop); self.jump_to_fail_pop(pc, JumpOp::Jump)?; @@ -6691,6 +6753,7 @@ impl Compiler { let n_rots = n_stores + 1 + pc.on_top + pc.stores.len(); for i in 0..n_stores { // Rotate the capture to its proper place. + self.set_source_range(p.range()); self.pattern_helper_rotate(n_rots)?; let name = &control.as_ref().unwrap()[i]; // Check for duplicate binding. @@ -6702,6 +6765,7 @@ impl Compiler { // Old context and control will be dropped automatically. // Finally, pop the copy of the subject. + self.set_source_range(p.range()); emit!(self, Instruction::PopTop); Ok(()) } @@ -6790,7 +6854,9 @@ impl Compiler { p: &ast::PatternMatchValue, pc: &mut PatternContext, ) -> CompileResult<()> { - // TODO: ensure literal or attribute lookup + // Match CPython codegen_pattern_value(): compare, then normalize to bool + // before the fail jump. Late IR folding will collapse COMPARE_OP+TO_BOOL + // into COMPARE_OP bool(...) when applicable. self.compile_expression(&p.value)?; emit!( self, @@ -6798,7 +6864,7 @@ impl Compiler { opname: bytecode::ComparisonOperator::Equal } ); - // emit!(self, Instruction::ToBool); + emit!(self, Instruction::ToBool); self.jump_to_fail_pop(pc, JumpOp::PopJumpIfFalse)?; Ok(()) } @@ -6826,7 +6892,9 @@ impl Compiler { pattern_type: &ast::Pattern, pattern_context: &mut PatternContext, ) -> CompileResult<()> { - match &pattern_type { + let prev_source_range = self.current_source_range; + self.set_source_range(pattern_type.range()); + let result = match &pattern_type { ast::Pattern::MatchValue(pattern_type) => { self.compile_pattern_value(pattern_type, pattern_context) } @@ -6851,7 +6919,9 @@ impl Compiler { ast::Pattern::MatchOr(pattern_type) => { self.compile_pattern_or(pattern_type, pattern_context) } - } + }; + self.set_source_range(prev_source_range); + result } fn compile_match_inner( @@ -6860,12 +6930,22 @@ impl Compiler { cases: &[ast::MatchCase], pattern_context: &mut PatternContext, ) -> CompileResult<()> { + fn is_trailing_wildcard_default(pattern: &ast::Pattern) -> bool { + match pattern { + ast::Pattern::MatchAs(match_as) => { + match_as.pattern.is_none() && match_as.name.is_none() + } + _ => false, + } + } + self.compile_expression(subject)?; let end = self.new_block(); let num_cases = cases.len(); assert!(num_cases > 0); - let has_default = cases.iter().last().unwrap().pattern.is_match_star() && num_cases > 1; + let has_default = + num_cases > 1 && is_trailing_wildcard_default(&cases.last().unwrap().pattern); let case_count = num_cases - if has_default { 1 } else { 0 }; for (i, m) in cases.iter().enumerate().take(case_count) { @@ -6875,36 +6955,41 @@ impl Compiler { } pattern_context.stores = Vec::with_capacity(1); - pattern_context.allow_irrefutable = m.guard.is_some() || i == case_count - 1; + pattern_context.allow_irrefutable = m.guard.is_some() || i == num_cases - 1; pattern_context.fail_pop.clear(); pattern_context.on_top = 0; self.compile_pattern(&m.pattern, pattern_context)?; assert_eq!(pattern_context.on_top, 0); + self.set_source_range(m.pattern.range()); for name in &pattern_context.stores { self.compile_name(name, NameUsage::Store)?; } if let Some(ref guard) = m.guard { self.ensure_fail_pop(pattern_context, 0)?; - // Compile the guard expression - self.compile_expression(guard)?; - emit!(self, Instruction::ToBool); - emit!( - self, - Instruction::PopJumpIfFalse { - delta: pattern_context.fail_pop[0] - } - ); + self.compile_jump_if_inner( + guard, + false, + pattern_context.fail_pop[0], + Some(m.pattern.range()), + )?; } if i != case_count - 1 { + if let Some(first_stmt) = m.body.first() { + self.set_source_range(first_stmt.range()); + } + if matches!(m.pattern, ast::Pattern::MatchOr(_)) { + emit!(self, Instruction::Nop); + } emit!(self, Instruction::PopTop); } self.compile_statements(&m.body)?; emit!(self, PseudoInstruction::Jump { delta: end }); + self.set_source_range(m.pattern.range()); self.emit_and_reset_fail_pop(pattern_context)?; } @@ -6912,15 +6997,11 @@ impl Compiler { let m = &cases[num_cases - 1]; if num_cases == 1 { emit!(self, Instruction::PopTop); - } else { + } else if m.guard.is_none() { emit!(self, Instruction::Nop); } if let Some(ref guard) = m.guard { - // Compile guard and jump to end if false - self.compile_expression(guard)?; - emit!(self, Instruction::Copy { i: 1 }); - emit!(self, Instruction::PopJumpIfFalse { delta: end }); - emit!(self, Instruction::PopTop); + self.compile_jump_if(guard, false, end)?; } self.compile_statements(&m.body)?; } @@ -7034,6 +7115,7 @@ impl Compiler { // if comparison result is false, we break with this value; if true, try the next one. emit!(self, Instruction::Copy { i: 1 }); + emit!(self, Instruction::ToBool); emit!(self, Instruction::PopJumpIfFalse { delta: cleanup }); emit!(self, Instruction::PopTop); } @@ -7085,12 +7167,14 @@ impl Compiler { emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::Copy { i: 2 }); self.compile_addcompare(op); + emit!(self, Instruction::ToBool); emit!(self, Instruction::PopJumpIfFalse { delta: cleanup }); } self.compile_expression(last_comparator)?; self.set_source_range(compare_range); self.compile_addcompare(last_op); + emit!(self, Instruction::ToBool); self.emit_pop_jump_by_condition(condition, target_block); emit!(self, PseudoInstruction::Jump { delta: end }); @@ -7156,6 +7240,37 @@ impl Compiler { Ok(()) } + fn compile_check_annotation_expression(&mut self, expression: &ast::Expr) -> CompileResult<()> { + self.compile_expression(expression)?; + emit!(self, Instruction::PopTop); + Ok(()) + } + + fn compile_check_annotation_subscript(&mut self, expression: &ast::Expr) -> CompileResult<()> { + match expression { + ast::Expr::Slice(ast::ExprSlice { + lower, upper, step, .. + }) => { + if let Some(lower) = lower { + self.compile_check_annotation_expression(lower)?; + } + if let Some(upper) = upper { + self.compile_check_annotation_expression(upper)?; + } + if let Some(step) = step { + self.compile_check_annotation_expression(step)?; + } + } + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { + for element in elts { + self.compile_check_annotation_subscript(element)?; + } + } + _ => self.compile_check_annotation_expression(expression)?, + } + Ok(()) + } + fn compile_annotated_assign( &mut self, target: &ast::Expr, @@ -7222,6 +7337,19 @@ impl Compiler { } } + if value.is_none() { + match target { + ast::Expr::Attribute(ast::ExprAttribute { value, .. }) => { + self.compile_check_annotation_expression(value)?; + } + ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => { + self.compile_check_annotation_expression(value)?; + self.compile_check_annotation_subscript(slice)?; + } + _ => {} + } + } + Ok(()) } @@ -7398,14 +7526,15 @@ impl Compiler { /// /// The idea is to jump to a label if the expression is either true or false /// (indicated by the condition parameter). - fn compile_jump_if( + fn compile_jump_if_inner( &mut self, expression: &ast::Expr, condition: bool, target_block: BlockIdx, + source_range: Option, ) -> CompileResult<()> { let prev_source_range = self.current_source_range; - self.set_source_range(expression.range()); + self.set_source_range(source_range.unwrap_or_else(|| expression.range())); // Compile expression for test, and jump to label if false let result = match &expression { @@ -7419,16 +7548,26 @@ impl Compiler { // If any of the values is false, we can short-circuit. for value in values { - self.compile_jump_if(value, false, end_block)?; + self.compile_jump_if_inner(value, false, end_block, source_range)?; } // It depends upon the last value now: will it be true? - self.compile_jump_if(last_value, true, target_block)?; + self.compile_jump_if_inner( + last_value, + true, + target_block, + source_range, + )?; self.switch_to_block(end_block); } else { // If any value is false, the whole condition is false. for value in values { - self.compile_jump_if(value, false, target_block)?; + self.compile_jump_if_inner( + value, + false, + target_block, + source_range, + )?; } } } @@ -7436,7 +7575,12 @@ impl Compiler { if condition { // If any of the values is true. for value in values { - self.compile_jump_if(value, true, target_block)?; + self.compile_jump_if_inner( + value, + true, + target_block, + source_range, + )?; } } else { // If all of the values are false. @@ -7445,11 +7589,16 @@ impl Compiler { // If any value is true, we can short-circuit: for value in values { - self.compile_jump_if(value, true, end_block)?; + self.compile_jump_if_inner(value, true, end_block, source_range)?; } // It all depends upon the last value now! - self.compile_jump_if(last_value, false, target_block)?; + self.compile_jump_if_inner( + last_value, + false, + target_block, + source_range, + )?; self.switch_to_block(end_block); } } @@ -7460,7 +7609,7 @@ impl Compiler { op: ast::UnaryOp::Not, operand, .. - }) => self.compile_jump_if(operand, !condition, target_block), + }) => self.compile_jump_if_inner(operand, !condition, target_block, source_range), ast::Expr::Compare(ast::ExprCompare { left, ops, @@ -7507,10 +7656,7 @@ impl Compiler { _ => { // Fall back case which always will work! self.compile_expression(expression)?; - // Compare already produces a bool; everything else needs TO_BOOL - if !matches!(expression, ast::Expr::Compare(_)) { - emit!(self, Instruction::ToBool); - } + emit!(self, Instruction::ToBool); if condition { emit!( self, @@ -7534,6 +7680,15 @@ impl Compiler { result } + fn compile_jump_if( + &mut self, + expression: &ast::Expr, + condition: bool, + target_block: BlockIdx, + ) -> CompileResult<()> { + self.compile_jump_if_inner(expression, condition, target_block, None) + } + /// Compile a boolean operation as an expression. /// This means, that the last value remains on the stack. fn compile_bool_op(&mut self, op: &ast::BoolOp, values: &[ast::Expr]) -> CompileResult<()> { @@ -7846,6 +8001,50 @@ impl Compiler { let range = expression.range(); self.set_source_range(range); + if !self.disable_const_boolop_folding + && let ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) = expression + { + let mut simplified_prefix = 0usize; + let mut last_constant = None; + for value in values { + let Some(constant) = self.try_fold_constant_expr(value)? else { + break; + }; + let is_truthy = Self::constant_truthiness(&constant); + last_constant = Some(constant); + match op { + ast::BoolOp::Or if is_truthy => { + self.emit_load_const(last_constant.expect("missing boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); + return Ok(()); + } + ast::BoolOp::And if !is_truthy => { + self.emit_load_const(last_constant.expect("missing boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); + return Ok(()); + } + ast::BoolOp::Or | ast::BoolOp::And => { + simplified_prefix += 1; + } + } + } + + if simplified_prefix == values.len() { + self.emit_load_const(last_constant.expect("missing folded boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); + return Ok(()); + } + if simplified_prefix > 0 { + let tail = &values[simplified_prefix..]; + if let [value] = tail { + self.compile_expression(value)?; + } else { + self.compile_bool_op(op, tail)?; + } + return Ok(()); + } + } + match &expression { ast::Expr::Call(ast::ExprCall { func, arguments, .. @@ -8381,8 +8580,6 @@ impl Compiler { } match id.as_str() { "tuple" => Some(BuiltinGeneratorCallKind::Tuple), - "list" => Some(BuiltinGeneratorCallKind::List), - "set" => Some(BuiltinGeneratorCallKind::Set), "all" => Some(BuiltinGeneratorCallKind::All), "any" => Some(BuiltinGeneratorCallKind::Any), _ => None, @@ -8391,34 +8588,27 @@ impl Compiler { /// Emit the optimized inline loop for builtin(genexpr) calls. /// - /// Stack on entry: `[func, iter]` where `iter` is the already-compiled - /// generator iterator and `func` is the builtin candidate. - /// On return the compiler is positioned at the fallback block with - /// `[func, iter]` still on the stack (for the normal CALL path). + /// Stack on entry: `[func]` where `func` is the builtin candidate. + /// On return the compiler is positioned at the fallback block so the + /// normal call path can compile the original generator argument again. fn optimize_builtin_generator_call( &mut self, kind: BuiltinGeneratorCallKind, + generator_expr: &ast::Expr, end: BlockIdx, ) -> CompileResult<()> { let common_constant = match kind { BuiltinGeneratorCallKind::Tuple => bytecode::CommonConstant::BuiltinTuple, - BuiltinGeneratorCallKind::List => bytecode::CommonConstant::BuiltinList, - BuiltinGeneratorCallKind::Set => bytecode::CommonConstant::BuiltinSet, BuiltinGeneratorCallKind::All => bytecode::CommonConstant::BuiltinAll, BuiltinGeneratorCallKind::Any => bytecode::CommonConstant::BuiltinAny, }; + let fallback = self.new_block(); let loop_block = self.new_block(); let cleanup = self.new_block(); - let fallback = self.new_block(); - let result = matches!( - kind, - BuiltinGeneratorCallKind::All | BuiltinGeneratorCallKind::Any - ) - .then(|| self.new_block()); - // Stack: [func, iter] — copy func (TOS1) for identity check - emit!(self, Instruction::Copy { i: 2 }); + // Stack: [func] — copy function for identity check + emit!(self, Instruction::Copy { i: 1 }); emit!( self, Instruction::LoadCommonConstant { @@ -8427,61 +8617,43 @@ impl Compiler { ); emit!(self, Instruction::IsOp { invert: Invert::No }); emit!(self, Instruction::PopJumpIfFalse { delta: fallback }); - emit!(self, Instruction::NotTaken); - // Remove func from [func, iter] → [iter] - emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::PopTop); - if matches!( - kind, - BuiltinGeneratorCallKind::Tuple | BuiltinGeneratorCallKind::List - ) { - // [iter] → [iter, list] → [list, iter] + if matches!(kind, BuiltinGeneratorCallKind::Tuple) { emit!(self, Instruction::BuildList { count: 0 }); - emit!(self, Instruction::Swap { i: 2 }); - } else if matches!(kind, BuiltinGeneratorCallKind::Set) { - // [iter] → [iter, set] → [set, iter] - emit!(self, Instruction::BuildSet { count: 0 }); - emit!(self, Instruction::Swap { i: 2 }); } + let sub_table_cursor = self.symbol_table_stack.last().map(|t| t.next_sub_table); + self.compile_expression(generator_expr)?; + if let Some(cursor) = sub_table_cursor + && let Some(current_table) = self.symbol_table_stack.last_mut() + { + current_table.next_sub_table = cursor; + } self.switch_to_block(loop_block); emit!(self, Instruction::ForIter { delta: cleanup }); match kind { - BuiltinGeneratorCallKind::Tuple | BuiltinGeneratorCallKind::List => { + BuiltinGeneratorCallKind::Tuple => { emit!(self, Instruction::ListAppend { i: 2 }); emit!(self, PseudoInstruction::Jump { delta: loop_block }); } - BuiltinGeneratorCallKind::Set => { - emit!(self, Instruction::SetAdd { i: 2 }); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); - } BuiltinGeneratorCallKind::All => { - let result = result.expect("all() optimization should have a result block"); emit!(self, Instruction::ToBool); - emit!(self, Instruction::PopJumpIfFalse { delta: result }); - emit!(self, Instruction::NotTaken); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + emit!(self, Instruction::PopJumpIfTrue { delta: loop_block }); + emit!(self, Instruction::PopIter); + self.emit_load_const(ConstantData::Boolean { value: false }); + emit!(self, PseudoInstruction::Jump { delta: end }); } BuiltinGeneratorCallKind::Any => { - let result = result.expect("any() optimization should have a result block"); emit!(self, Instruction::ToBool); - emit!(self, Instruction::PopJumpIfTrue { delta: result }); - emit!(self, Instruction::NotTaken); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + emit!(self, Instruction::PopJumpIfFalse { delta: loop_block }); + emit!(self, Instruction::PopIter); + self.emit_load_const(ConstantData::Boolean { value: true }); + emit!(self, PseudoInstruction::Jump { delta: end }); } } - if let Some(result_block) = result { - self.switch_to_block(result_block); - emit!(self, Instruction::PopIter); - self.emit_load_const(ConstantData::Boolean { - value: matches!(kind, BuiltinGeneratorCallKind::Any), - }); - emit!(self, PseudoInstruction::Jump { delta: end }); - } - self.switch_to_block(cleanup); emit!(self, Instruction::EndFor); emit!(self, Instruction::PopIter); @@ -8494,7 +8666,6 @@ impl Compiler { } ); } - BuiltinGeneratorCallKind::List | BuiltinGeneratorCallKind::Set => {} BuiltinGeneratorCallKind::All => { self.emit_load_const(ConstantData::Boolean { value: true }); } @@ -8574,18 +8745,12 @@ impl Compiler { .then(|| self.detect_builtin_generator_call(func, args)) .flatten() { - // Optimized builtin(genexpr) path: compile the genexpr only once - // so its code object appears exactly once in co_consts. let end = self.new_block(); self.compile_expression(func)?; - self.compile_expression(&args.args[0])?; - // Stack: [func, iter] - self.optimize_builtin_generator_call(kind, end)?; - // Fallback block: [func, iter] → [func, null, iter] → CALL - emit!(self, Instruction::PushNull); - emit!(self, Instruction::Swap { i: 2 }); + self.optimize_builtin_generator_call(kind, &args.args[0], end)?; self.set_source_range(call_range); - emit!(self, Instruction::Call { argc: 1 }); + emit!(self, Instruction::PushNull); + self.codegen_call_helper(0, args, call_range)?; self.switch_to_block(end); } else { // Regular call: push func, then NULL for self_or_null slot @@ -8707,7 +8872,7 @@ impl Compiler { // Single starred arg: pass value directly to CallFunctionEx. // Runtime will convert to tuple and validate with function name. if let ast::Expr::Starred(ast::ExprStarred { value, .. }) = &arguments.args[0] { - self.compile_expression(value)?; + self.compile_expression_without_const_boolop_folding(value)?; } } else if !has_starred { for arg in &arguments.args { @@ -8763,7 +8928,7 @@ impl Compiler { have_dict = true; } - self.compile_expression(&keyword.value)?; + self.compile_expression_without_const_boolop_folding(&keyword.value)?; emit!(self, Instruction::DictMerge { i: 1 }); } else { nseen += 1; @@ -8857,20 +9022,16 @@ impl Compiler { ast::Expr::ListComp(ast::ExprListComp { generators, .. }) | ast::Expr::SetComp(ast::ExprSetComp { generators, .. }) | ast::Expr::Generator(ast::ExprGenerator { generators, .. }) => { - // leave_scope runs before the first iterator is - // scanned, so the comprehension scope comes first - // in sub_tables, then any nested scopes from the - // first iterator. - self.consume_scope(); if let Some(first) = generators.first() { self.visit_expr(&first.iter); } + self.consume_scope(); } ast::Expr::DictComp(ast::ExprDictComp { generators, .. }) => { - self.consume_scope(); if let Some(first) = generators.first() { self.visit_expr(&first.iter); } + self.consume_scope(); } _ => ast::visitor::walk_expr(self, expr), } @@ -8889,30 +9050,85 @@ impl Compiler { } } - fn compile_comprehension( + fn peek_next_sub_table_after_skipped_nested_scopes_in_expr( &mut self, - name: &str, - init_collection: Option, - generators: &[ast::Comprehension], - compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, - comprehension_type: ComprehensionType, - element_contains_await: bool, - ) -> CompileResult<()> { - let prev_ctx = self.ctx; - let has_an_async_gen = generators.iter().any(|g| g.is_async); - - // Check for async comprehension outside async function (list/set/dict only, not generator expressions) - // Use in_async_scope to allow nested async comprehensions inside an async function - if comprehension_type != ComprehensionType::Generator - && (has_an_async_gen || element_contains_await) - && !prev_ctx.in_async_scope + expression: &ast::Expr, + ) -> CompileResult { + let saved_cursor = self + .symbol_table_stack + .last() + .expect("no current symbol table") + .next_sub_table; + let result = (|| { + self.consume_skipped_nested_scopes_in_expr(expression)?; + let current_table = self + .symbol_table_stack + .last() + .expect("no current symbol table"); + if let Some(table) = current_table.sub_tables.get(current_table.next_sub_table) { + Ok(table.clone()) + } else { + let name = current_table.name.clone(); + let typ = current_table.typ; + Err(self.error(CodegenErrorType::SyntaxError(format!( + "no symbol table available in {} (type: {:?})", + name, typ + )))) + } + })(); + self.symbol_table_stack + .last_mut() + .expect("no current symbol table") + .next_sub_table = saved_cursor; + result + } + + fn push_output_with_symbol_table( + &mut self, + table: SymbolTable, + flags: bytecode::CodeFlags, + posonlyarg_count: u32, + arg_count: u32, + kwonlyarg_count: u32, + obj_name: String, + ) -> CompileResult<()> { + let scope_type = table.typ; + self.symbol_table_stack.push(table); + + let key = self.symbol_table_stack.len() - 1; + let lineno = self.get_source_line_number().get(); + self.enter_scope(&obj_name, scope_type, key, lineno.to_u32())?; + + if let Some(info) = self.code_stack.last_mut() { + info.flags = flags | (info.flags & bytecode::CodeFlags::NESTED); + info.metadata.argcount = arg_count; + info.metadata.posonlyargcount = posonlyarg_count; + info.metadata.kwonlyargcount = kwonlyarg_count; + } + Ok(()) + } + + fn compile_comprehension( + &mut self, + name: &str, + init_collection: Option, + generators: &[ast::Comprehension], + compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, + comprehension_type: ComprehensionType, + element_contains_await: bool, + ) -> CompileResult<()> { + let prev_ctx = self.ctx; + let has_an_async_gen = generators.iter().any(|g| g.is_async); + + // Check for async comprehension outside async function (list/set/dict only, not generator expressions) + // Use in_async_scope to allow nested async comprehensions inside an async function + if comprehension_type != ComprehensionType::Generator + && (has_an_async_gen || element_contains_await) + && !prev_ctx.in_async_scope { return Err(self.error(CodegenErrorType::InvalidAsyncComprehension)); } - // Check if this comprehension should be inlined (PEP 709) - let is_inlined = self.is_inlined_comprehension_context(comprehension_type); - // async comprehensions are allowed in various contexts: // - list/set/dict comprehensions in async functions (or nested within) // - always for generator expressions @@ -8930,12 +9146,18 @@ impl Compiler { // We must have at least one generator: assert!(!generators.is_empty()); + let outermost = &generators[0]; + let comp_table = + self.peek_next_sub_table_after_skipped_nested_scopes_in_expr(&outermost.iter)?; + + let is_inlined = self.is_inlined_comprehension_context(comprehension_type, &comp_table); if is_inlined && !has_an_async_gen && !element_contains_await { // PEP 709: Inlined comprehension - compile inline without new scope let was_in_inlined_comp = self.current_code_info().in_inlined_comp; self.current_code_info().in_inlined_comp = true; let result = self.compile_inlined_comprehension( + comp_table, init_collection, generators, compile_element, @@ -8966,8 +9188,12 @@ impl Compiler { flags }; - // Create magnificent function : - self.push_output(flags, 1, 1, 0, name.to_owned())?; + // The symbol table follows CPython's symtable walk: nested scopes + // in the outermost iterator are recorded before the comprehension + // scope itself. Peek past those nested scopes so we can enter the + // correct comprehension table here, then let the real outermost + // iterator compile consume its nested scopes later in parent scope. + self.push_output_with_symbol_table(comp_table, flags, 1, 1, 0, name.to_owned())?; // Set qualname for comprehension self.set_qualname(); @@ -9009,7 +9235,7 @@ impl Compiler { emit!(self, Instruction::LoadFast { var_num: arg0 }); } else { // Evaluate iterated item: - self.compile_expression(&generator.iter)?; + self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; // Get iterator / turn item into an iterator if generator.is_async { @@ -9048,9 +9274,14 @@ impl Compiler { end_async_for_target, )); - // Now evaluate the ifs: + // CPython always lowers comprehension guards through codegen_jump_if + // and leaves constant-folding to later CFG optimization passes. for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, if_cleanup_block)? + self.compile_jump_if(if_condition, false, if_cleanup_block)?; + } + if !generator.ifs.is_empty() { + let body_block = self.new_block(); + self.switch_to_block(body_block); } } @@ -9107,11 +9338,15 @@ impl Compiler { self.make_closure(code, bytecode::MakeFunctionFlags::new())?; // Evaluate iterated item: - self.compile_expression(&generators[0].iter)?; + self.compile_for_iterable_expression(&outermost.iter, outermost.is_async)?; + self.symbol_table_stack + .last_mut() + .expect("no current symbol table") + .next_sub_table += 1; // Get iterator / turn item into an iterator // Use is_async from the first generator, not has_an_async_gen which covers ALL generators - if generators[0].is_async { + if outermost.is_async { emit!(self, Instruction::GetAIter); } else { emit!(self, Instruction::GetIter); @@ -9132,25 +9367,24 @@ impl Compiler { /// This generates bytecode inline without creating a new code object fn compile_inlined_comprehension( &mut self, + comp_table: SymbolTable, init_collection: Option, generators: &[ast::Comprehension], compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, has_async: bool, ) -> CompileResult<()> { - // PEP 709: Consume the comprehension's sub_table. - // The symbols are already merged into parent scope by analyze_symbol_table. - let current_table = self - .symbol_table_stack - .last_mut() - .expect("no current symbol table"); - let comp_table = current_table.sub_tables[current_table.next_sub_table].clone(); - current_table.next_sub_table += 1; - // Compile the outermost iterator first. Its expression may reference // nested scopes (e.g. lambdas) whose sub_tables sit at the current // position in the parent's list. Those must be consumed before we // splice in the comprehension's own children. - self.compile_expression(&generators[0].iter)?; + self.compile_for_iterable_expression( + &generators[0].iter, + has_async && generators[0].is_async, + )?; + self.symbol_table_stack + .last_mut() + .expect("no current symbol table") + .next_sub_table += 1; // Splice the comprehension's children (e.g. nested inlined // comprehensions) into the parent so the compiler can find them. @@ -9176,22 +9410,49 @@ impl Compiler { let ct = self.current_symbol_table(); ct.typ == CompilerScope::Class && !self.current_code_info().in_inlined_comp }; + fn collect_bound_names(target: &ast::Expr, out: &mut Vec) { + match target { + ast::Expr::Name(ast::ExprName { id, .. }) => out.push(id.to_string()), + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) + | ast::Expr::List(ast::ExprList { elts, .. }) => { + for elt in elts { + collect_bound_names(elt, out); + } + } + ast::Expr::Starred(ast::ExprStarred { value, .. }) => { + collect_bound_names(value, out); + } + _ => {} + } + } + let mut source_order_bound_names = Vec::new(); + for generator in generators { + collect_bound_names(&generator.target, &mut source_order_bound_names); + } let mut pushed_locals: Vec = Vec::new(); - for (name, sym) in &comp_table.symbols { - if sym.flags.contains(SymbolFlags::PARAMETER) { - continue; // skip .0 + for name in source_order_bound_names + .into_iter() + .chain(comp_table.symbols.keys().cloned()) + { + if pushed_locals.iter().any(|existing| existing == &name) { + continue; } - // Walrus operator targets (ASSIGNED_IN_COMPREHENSION without ITER) - // are not local to the comprehension; they leak to the outer scope. - let is_walrus = sym.flags.contains(SymbolFlags::ASSIGNED_IN_COMPREHENSION) - && !sym.flags.contains(SymbolFlags::ITER); - let is_local = sym - .flags - .intersects(SymbolFlags::ASSIGNED | SymbolFlags::ITER) - && !sym.flags.contains(SymbolFlags::NONLOCAL) - && !is_walrus; - if is_local || in_class_block { - pushed_locals.push(name.clone()); + if let Some(sym) = comp_table.symbols.get(&name) { + if sym.flags.contains(SymbolFlags::PARAMETER) { + continue; // skip .0 + } + // Walrus operator targets (ASSIGNED_IN_COMPREHENSION without ITER) + // are not local to the comprehension; they leak to the outer scope. + let is_walrus = sym.flags.contains(SymbolFlags::ASSIGNED_IN_COMPREHENSION) + && !sym.flags.contains(SymbolFlags::ITER); + let is_local = sym + .flags + .intersects(SymbolFlags::ASSIGNED | SymbolFlags::ITER) + && !sym.flags.contains(SymbolFlags::NONLOCAL) + && !is_walrus; + if is_local || in_class_block { + pushed_locals.push(name); + } } } @@ -9287,7 +9548,7 @@ impl Compiler { let after_block = self.new_block(); if i > 0 { - self.compile_expression(&generator.iter)?; + self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; if generator.is_async { emit!(self, Instruction::GetAIter); } else { @@ -9324,7 +9585,8 @@ impl Compiler { end_async_for_target, )); - // Evaluate the if conditions + // CPython always lowers comprehension guards through codegen_jump_if + // and leaves constant-folding to later CFG optimization passes. for if_condition in &generator.ifs { self.compile_jump_if(if_condition, false, if_cleanup_block)?; } @@ -9446,11 +9708,18 @@ impl Compiler { location, end_location, except_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); } + fn mark_last_instruction_folded_from_nonliteral_expr(&mut self) { + if let Some(info) = self.current_block().instructions.last_mut() { + info.folded_from_nonliteral_expr = true; + } + } + /// Mark the last emitted instruction as having no source location. /// Prevents it from triggering LINE events in sys.monitoring. fn set_no_location(&mut self) { @@ -9505,6 +9774,19 @@ impl Compiler { } } + fn compile_tstring_literal_value( + &self, + string: &ast::InterpolatedStringLiteralElement, + flags: ast::TStringFlags, + ) -> Wtf8Buf { + if string.value.contains(char::REPLACEMENT_CHARACTER) { + let source = self.source_file.slice(string.range); + crate::string_parser::parse_fstring_literal_element(source.into(), flags.into()).into() + } else { + string.value.to_string().into() + } + } + fn compile_fstring_part_literal_value(&self, string: &ast::StringLiteral) -> Wtf8Buf { if string.value.contains(char::REPLACEMENT_CHARACTER) { let source = self.source_file.slice(string.range); @@ -9516,14 +9798,85 @@ impl Compiler { fn arg_constant(&mut self, constant: ConstantData) -> oparg::ConstIdx { let info = self.current_code_info(); + if let ConstantData::Code { code } = &constant + && let Some(idx) = info.metadata.consts.iter().position(|existing| { + matches!( + existing, + ConstantData::Code { + code: existing_code + } if Self::code_objects_equivalent(existing_code, code) + ) + }) + { + return u32::try_from(idx) + .expect("constant table index overflow") + .into(); + } info.metadata.consts.insert_full(constant).0.to_u32().into() } + fn constants_equivalent(lhs: &ConstantData, rhs: &ConstantData) -> bool { + match (lhs, rhs) { + (ConstantData::Code { code: lhs }, ConstantData::Code { code: rhs }) => { + Self::code_objects_equivalent(lhs, rhs) + } + (ConstantData::Tuple { elements: lhs }, ConstantData::Tuple { elements: rhs }) + | ( + ConstantData::Frozenset { elements: lhs }, + ConstantData::Frozenset { elements: rhs }, + ) => { + lhs.len() == rhs.len() + && lhs + .iter() + .zip(rhs.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)) + } + (ConstantData::Slice { elements: lhs }, ConstantData::Slice { elements: rhs }) => lhs + .iter() + .zip(rhs.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)), + _ => lhs == rhs, + } + } + + fn code_objects_equivalent(lhs: &bytecode::CodeObject, rhs: &bytecode::CodeObject) -> bool { + lhs.instructions.len() == rhs.instructions.len() + && lhs + .instructions + .iter() + .zip(rhs.instructions.iter()) + .all(|(lhs, rhs)| u8::from(lhs.op) == u8::from(rhs.op) && lhs.arg == rhs.arg) + && lhs.locations == rhs.locations + && lhs.flags.bits() == rhs.flags.bits() + && lhs.posonlyarg_count == rhs.posonlyarg_count + && lhs.arg_count == rhs.arg_count + && lhs.kwonlyarg_count == rhs.kwonlyarg_count + && lhs.source_path == rhs.source_path + && lhs.first_line_number == rhs.first_line_number + && lhs.max_stackdepth == rhs.max_stackdepth + && lhs.obj_name == rhs.obj_name + && lhs.qualname == rhs.qualname + && lhs.constants.len() == rhs.constants.len() + && lhs + .constants + .iter() + .zip(rhs.constants.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)) + && lhs.names == rhs.names + && lhs.varnames == rhs.varnames + && lhs.cellvars == rhs.cellvars + && lhs.freevars == rhs.freevars + && lhs.localspluskinds == rhs.localspluskinds + && lhs.linetable == rhs.linetable + && lhs.exceptiontable == rhs.exceptiontable + } + /// Try to fold a collection of constant expressions into a single ConstantData::Tuple. /// Returns None if any element cannot be folded. fn try_fold_constant_collection( &mut self, elts: &[ast::Expr], + collection_type: CollectionType, ) -> CompileResult> { let mut constants = Vec::with_capacity(elts.len()); for elt in elts { @@ -9532,9 +9885,15 @@ impl Compiler { }; constants.push(constant); } - Ok(Some(ConstantData::Tuple { - elements: constants, - })) + let constant = match collection_type { + CollectionType::Tuple | CollectionType::List => ConstantData::Tuple { + elements: constants, + }, + CollectionType::Set => ConstantData::Frozenset { + elements: constants, + }, + }; + Ok(Some(constant)) } fn try_fold_constant_expr(&mut self, expr: &ast::Expr) -> CompileResult> { @@ -9567,6 +9926,132 @@ impl Compiler { } ConstantData::Tuple { elements } } + ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => { + let Some(container) = self.try_fold_constant_expr(value)? else { + return Ok(None); + }; + let Some(index) = self.try_fold_constant_expr(slice)? else { + return Ok(None); + }; + let ConstantData::Integer { value: index } = index else { + return Ok(None); + }; + let Some(index): Option = index.try_into().ok() else { + return Ok(None); + }; + + match container { + ConstantData::Str { value } => { + let string = value.to_string(); + if string.contains(char::REPLACEMENT_CHARACTER) { + return Ok(None); + } + let chars: Vec<_> = string.chars().collect(); + let Some(len) = i64::try_from(chars.len()).ok() else { + return Ok(None); + }; + let idx: i64 = if index < 0 { len + index } else { index }; + let Some(idx) = usize::try_from(idx).ok() else { + return Ok(None); + }; + let Some(ch) = chars.get(idx) else { + return Ok(None); + }; + ConstantData::Str { + value: ch.to_string().into(), + } + } + ConstantData::Bytes { value } => { + let Some(len) = i64::try_from(value.len()).ok() else { + return Ok(None); + }; + let idx: i64 = if index < 0 { len + index } else { index }; + let Some(idx) = usize::try_from(idx).ok() else { + return Ok(None); + }; + let Some(byte) = value.get(idx) else { + return Ok(None); + }; + ConstantData::Integer { + value: BigInt::from(*byte), + } + } + ConstantData::Tuple { elements } => { + let Some(len) = i64::try_from(elements.len()).ok() else { + return Ok(None); + }; + let idx: i64 = if index < 0 { len + index } else { index }; + let Some(idx) = usize::try_from(idx).ok() else { + return Ok(None); + }; + let Some(element) = elements.get(idx) else { + return Ok(None); + }; + element.clone() + } + _ => return Ok(None), + } + } + ast::Expr::UnaryOp(ast::ExprUnaryOp { op, operand, .. }) => { + let Some(constant) = self.try_fold_constant_expr(operand)? else { + return Ok(None); + }; + match (op, constant) { + (ast::UnaryOp::UAdd, value) => value, + (ast::UnaryOp::USub, ConstantData::Integer { value }) => { + ConstantData::Integer { value: -value } + } + (ast::UnaryOp::USub, ConstantData::Float { value }) => { + ConstantData::Float { value: -value } + } + (ast::UnaryOp::USub, ConstantData::Complex { value }) => { + ConstantData::Complex { value: -value } + } + (ast::UnaryOp::Invert, ConstantData::Integer { value }) => { + ConstantData::Integer { value: !value } + } + _ => return Ok(None), + } + } + ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => { + let mut constants = Vec::with_capacity(values.len()); + for value in values { + let Some(constant) = self.try_fold_constant_expr(value)? else { + return Ok(None); + }; + constants.push(constant); + } + let mut iter = constants.into_iter(); + let Some(first) = iter.next() else { + return Ok(None); + }; + let mut selected = first; + match op { + ast::BoolOp::Or => { + if !Self::constant_truthiness(&selected) { + for constant in iter { + let is_truthy = Self::constant_truthiness(&constant); + selected = constant; + if is_truthy { + break; + } + } + } + } + ast::BoolOp::And => { + if Self::constant_truthiness(&selected) { + for constant in iter { + let is_truthy = Self::constant_truthiness(&constant); + selected = constant; + if !is_truthy { + break; + } + } + } + } + } + selected + } _ => return Ok(None), })) } @@ -9670,44 +10155,6 @@ impl Compiler { self.code_stack.last_mut().expect("no code on stack") } - /// Evaluate whether an expression is a compile-time constant boolean. - /// Returns Some(true) for truthy constants, Some(false) for falsy constants, - /// None for non-constant expressions. - /// = expr_constant in CPython compile.c - fn expr_constant(expr: &ast::Expr) -> Option { - match expr { - ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value), - ast::Expr::NoneLiteral(_) => Some(false), - ast::Expr::EllipsisLiteral(_) => Some(true), - ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => match value { - ast::Number::Int(i) => { - let n: i64 = i.as_i64().unwrap_or(1); - Some(n != 0) - } - ast::Number::Float(f) => Some(*f != 0.0), - ast::Number::Complex { real, imag, .. } => Some(*real != 0.0 || *imag != 0.0), - }, - ast::Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => { - Some(!value.to_str().is_empty()) - } - ast::Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => { - Some(value.bytes().next().is_some()) - } - ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { - if elts.is_empty() { - Some(false) - } else { - None // non-empty tuples may have side effects in elements - } - } - _ => None, - } - } - - fn emit_nop(&mut self) { - emit!(self, Instruction::Nop); - } - /// Enter a conditional block (if/for/while/match/try/with) /// PEP 649: Track conditional annotation context fn enter_conditional_block(&mut self) { @@ -9800,6 +10247,11 @@ impl Compiler { let loop_block = code.fblock[loop_idx].fb_block; let exit_block = code.fblock[loop_idx].fb_exit; + let prev_source_range = self.current_source_range; + self.set_source_range(range); + emit!(self, Instruction::Nop); + self.set_source_range(prev_source_range); + // Collect the fblocks we need to unwind through, from top down to (but not including) the loop #[derive(Clone)] enum UnwindAction { @@ -9949,7 +10401,13 @@ impl Compiler { fn new_block(&mut self) -> BlockIdx { let code = self.current_code_info(); let idx = BlockIdx::new(code.blocks.len().to_u32()); - code.blocks.push(ir::Block::default()); + let inherited_disable_load_fast_borrow = + code.blocks[code.current_block].disable_load_fast_borrow; + let block = ir::Block { + disable_load_fast_borrow: inherited_disable_load_fast_borrow, + ..ir::Block::default() + }; + code.blocks.push(block); idx } @@ -10342,39 +10800,26 @@ impl Compiler { } fn compile_expr_tstring(&mut self, expr_tstring: &ast::ExprTString) -> CompileResult<()> { - // ast::TStringValue can contain multiple ast::TString parts (implicit concatenation) - // Each ast::TString part should be compiled and the results merged into a single Template + // ast::TStringValue can contain multiple ast::TString parts (implicit + // concatenation). Match CPython's stack order by materializing the + // strings tuple first, then evaluating interpolations left-to-right. let tstring_value = &expr_tstring.value; - // Collect all strings and compile all interpolations let mut all_strings: Vec = Vec::new(); let mut current_string = Wtf8Buf::new(); let mut interp_count: u32 = 0; for tstring in tstring_value.iter() { - self.compile_tstring_into( + self.collect_tstring_strings( tstring, &mut all_strings, &mut current_string, &mut interp_count, - )?; + ); } - // Add trailing string all_strings.push(core::mem::take(&mut current_string)); - // Now build the Template: - // Stack currently has all interpolations from compile_tstring_into calls - - // 1. Build interpolations tuple from the interpolations on the stack - emit!( - self, - Instruction::BuildTuple { - count: interp_count - } - ); - - // 2. Load all string parts let string_count: u32 = all_strings .len() .try_into() @@ -10382,8 +10827,6 @@ impl Compiler { for s in &all_strings { self.emit_load_const(ConstantData::Str { value: s.clone() }); } - - // 3. Build strings tuple emit!( self, Instruction::BuildTuple { @@ -10391,79 +10834,95 @@ impl Compiler { } ); - // 4. Swap so strings is below interpolations: [interps, strings] -> [strings, interps] - emit!(self, Instruction::Swap { i: 2 }); + for tstring in tstring_value.iter() { + self.compile_tstring_interpolations(tstring)?; + } - // 5. Build the Template + emit!( + self, + Instruction::BuildTuple { + count: interp_count + } + ); emit!(self, Instruction::BuildTemplate); Ok(()) } - fn compile_tstring_into( - &mut self, + fn collect_tstring_strings( + &self, tstring: &ast::TString, strings: &mut Vec, current_string: &mut Wtf8Buf, interp_count: &mut u32, - ) -> CompileResult<()> { + ) { for element in &tstring.elements { match element { ast::InterpolatedStringElement::Literal(lit) => { - // Accumulate literal parts into current_string - current_string.push_str(&lit.value); + current_string + .push_wtf8(&self.compile_tstring_literal_value(lit, tstring.flags)); } ast::InterpolatedStringElement::Interpolation(interp) => { - // Finish current string segment + if let Some(ast::DebugText { leading, trailing }) = &interp.debug_text { + let range = interp.expression.range(); + let source = self.source_file.slice(range); + let text = [ + strip_fstring_debug_comments(leading).as_str(), + source, + strip_fstring_debug_comments(trailing).as_str(), + ] + .concat(); + current_string.push_str(&text); + } strings.push(core::mem::take(current_string)); + *interp_count += 1; + } + } + } + } - // Compile the interpolation value - self.compile_expression(&interp.expression)?; + fn compile_tstring_interpolations(&mut self, tstring: &ast::TString) -> CompileResult<()> { + for element in &tstring.elements { + let ast::InterpolatedStringElement::Interpolation(interp) = element else { + continue; + }; - // Load the expression source string, including any - // whitespace between '{' and the expression start - let expr_range = interp.expression.range(); - let expr_source = if interp.range.start() < expr_range.start() - && interp.range.end() >= expr_range.end() - { - let after_brace = interp.range.start() + TextSize::new(1); - self.source_file - .slice(TextRange::new(after_brace, expr_range.end())) - } else { - // Fallback for programmatically constructed ASTs with dummy ranges - self.source_file.slice(expr_range) - }; - self.emit_load_const(ConstantData::Str { - value: expr_source.to_string().into(), - }); + self.compile_expression(&interp.expression)?; - // Determine conversion code - let conversion: u32 = match interp.conversion { - ast::ConversionFlag::None => 0, - ast::ConversionFlag::Str => 1, - ast::ConversionFlag::Repr => 2, - ast::ConversionFlag::Ascii => 3, - }; + let expr_range = interp.expression.range(); + let expr_source = if interp.range.start() < expr_range.start() + && interp.range.end() >= expr_range.end() + { + let after_brace = interp.range.start() + TextSize::new(1); + self.source_file + .slice(TextRange::new(after_brace, expr_range.end())) + } else { + self.source_file.slice(expr_range) + }; + self.emit_load_const(ConstantData::Str { + value: expr_source.to_string().into(), + }); - // Handle format_spec - let has_format_spec = interp.format_spec.is_some(); - if let Some(format_spec) = &interp.format_spec { - // Compile format_spec as a string using fstring element compilation - // Use default ast::FStringFlags since format_spec syntax is independent of t-string flags - self.compile_fstring_elements( - ast::FStringFlags::empty(), - &format_spec.elements, - )?; - } + let mut conversion: u32 = match interp.conversion { + ast::ConversionFlag::None => 0, + ast::ConversionFlag::Str => 1, + ast::ConversionFlag::Repr => 2, + ast::ConversionFlag::Ascii => 3, + }; - // Emit BUILD_INTERPOLATION - // oparg encoding: (conversion << 2) | has_format_spec - let format = (conversion << 2) | u32::from(has_format_spec); - emit!(self, Instruction::BuildInterpolation { format }); + if interp.debug_text.is_some() && conversion == 0 && interp.format_spec.is_none() { + conversion = 2; + } - *interp_count += 1; - } + let has_format_spec = interp.format_spec.is_some(); + if let Some(format_spec) = &interp.format_spec { + self.compile_fstring_elements(ast::FStringFlags::empty(), &format_spec.elements)?; } + + // CPython keeps bit 1 set in BUILD_INTERPOLATION's oparg and uses + // bit 0 for the optional format spec. + let format = 2 | (conversion << 2) | u32::from(has_format_spec); + emit!(self, Instruction::BuildInterpolation { format }); } Ok(()) @@ -10840,18 +11299,103 @@ mod tests { compiler.exit_scope() } - fn find_code<'a>(code: &'a CodeObject, name: &str) -> Option<&'a CodeObject> { - if code.obj_name == name { - return Some(code); - } - code.constants.iter().find_map(|constant| { - if let ConstantData::Code { code } = constant { - find_code(code, name) - } else { - None - } - }) - } + fn compile_exec_late_cfg_trace(source: &str) -> Vec<(String, String)> { + let opts = CompileOpts::default(); + let source_file = SourceFileBuilder::new("source_path", source).finish(); + let parsed = ruff_python_parser::parse( + source_file.source_text(), + ruff_python_parser::Mode::Module.into(), + ) + .unwrap(); + let ast = parsed.into_syntax(); + let ast = match ast { + ruff_python_ast::Mod::Module(stmts) => stmts, + _ => unreachable!(), + }; + let symbol_table = SymbolTable::scan_program(&ast, source_file.clone()) + .map_err(|e| e.into_codegen_error(source_file.name().to_owned())) + .unwrap(); + let mut compiler = Compiler::new(opts, source_file, "".to_owned()); + compiler.compile_program(&ast, symbol_table).unwrap(); + let _table = compiler.pop_symbol_table(); + let stack_top = compiler.code_stack.pop().unwrap(); + stack_top.debug_late_cfg_trace().unwrap() + } + + fn compile_single_function_late_cfg_trace( + source: &str, + function_name: &str, + ) -> Vec<(String, String)> { + let opts = CompileOpts::default(); + let source_file = SourceFileBuilder::new("source_path", source).finish(); + let parsed = ruff_python_parser::parse( + source_file.source_text(), + ruff_python_parser::Mode::Module.into(), + ) + .unwrap(); + let ast = parsed.into_syntax(); + let ast = match ast { + ruff_python_ast::Mod::Module(stmts) => stmts, + _ => unreachable!(), + }; + let symbol_table = SymbolTable::scan_program(&ast, source_file.clone()) + .map_err(|e| e.into_codegen_error(source_file.name().to_owned())) + .unwrap(); + let function = ast + .body + .iter() + .find_map(|stmt| match stmt { + ast::Stmt::FunctionDef(f) if f.name.as_str() == function_name => Some(f), + _ => None, + }) + .unwrap_or_else(|| panic!("missing function {function_name}")); + + let mut compiler = Compiler::new(opts, source_file, "".to_owned()); + compiler.future_annotations = symbol_table.future_annotations; + compiler.symbol_table_stack.push(symbol_table); + compiler.set_source_range(function.range()); + compiler + .enter_function(function.name.as_str(), &function.parameters) + .unwrap(); + compiler + .current_code_info() + .flags + .set(bytecode::CodeFlags::COROUTINE, false); + + let prev_ctx = compiler.ctx; + compiler.ctx = CompileContext { + loop_data: None, + in_class: prev_ctx.in_class, + func: FunctionContext::Function, + in_async_scope: false, + }; + compiler.set_qualname(); + compiler.compile_statements(&function.body).unwrap(); + match function.body.last() { + Some(ast::Stmt::Return(_)) => {} + _ => compiler.emit_return_const(ConstantData::None), + } + if compiler.current_code_info().metadata.consts.is_empty() { + compiler.arg_constant(ConstantData::None); + } + + let _table = compiler.pop_symbol_table(); + let stack_top = compiler.code_stack.pop().unwrap(); + stack_top.debug_late_cfg_trace().unwrap() + } + + fn find_code<'a>(code: &'a CodeObject, name: &str) -> Option<&'a CodeObject> { + if code.obj_name == name { + return Some(code); + } + code.constants.iter().find_map(|constant| { + if let ConstantData::Code { code } = constant { + find_code(code, name) + } else { + None + } + }) + } fn has_common_constant(code: &CodeObject, expected: bytecode::CommonConstant) -> bool { code.instructions.iter().any(|unit| match unit.op { @@ -10891,6 +11435,195 @@ if True or False or False: )); } + #[test] + fn test_trace_assert_true_try_pair() { + let trace = compile_exec_late_cfg_trace( + "\ +try: + assert True +except AssertionError as e: + fail() +try: + assert True, 'msg' +except AssertionError as e: + fail() +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_for_unpack_list_literal() { + let trace = compile_exec_late_cfg_trace( + "\ +result = [] +for x, in [(1,), (2,), (3,)]: + result.append(x) +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_break_in_finally_function() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(self): + count = 0 + while count < 2: + count += 1 + try: + pass + finally: + break + self.assertEqual(count, 1) +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_constant_false_elif_chain() { + let trace = compile_exec_late_cfg_trace( + "\ +if 0: pass +elif 0: pass +elif 0: pass +elif 0: pass +else: pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_multi_pass_suite() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1: + # + # + # + pass + pass + # + pass + # +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_single_compare_if() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1 == 1: + pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_comparison_suite() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1: pass +x = (1 == 1) +if 1 == 1: pass +if 1 != 1: pass +if 1 < 1: pass +if 1 > 1: pass +if 1 <= 1: pass +if 1 >= 1: pass +if x is x: pass +if x is not x: pass +if 1 in (): pass +if 1 not in (): pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_if_for_except_layout() { + let trace = compile_exec_late_cfg_trace( + "\ +from sys import maxsize +if maxsize == 2147483647: + for s in ('2147483648', '0o40000000000', '0x100000000', '0b10000000000000000000000000000000'): + try: + x = eval(s) + except OverflowError: + fail(\"OverflowError on huge integer literal %r\" % s) +elif maxsize == 9223372036854775807: + pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_break_in_finally_tail_loads_borrow_through_empty_fallthrough_block() { + let code = compile_exec( + "\ +def f(self): + count = 0 + while count < 2: + count += 1 + try: + pass + finally: + break + self.assertEqual(count, 1) +", + ); + let code = find_code(&code, "f").unwrap(); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::LoadFastBorrow { .. }, + Instruction::LoadAttr { .. }, + Instruction::LoadFastBorrow { .. }, + Instruction::LoadSmallInt { .. }, + Instruction::Call { .. } + ] + ) + }), + "{:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + #[test] fn test_if_ands() { assert_dis_snapshot!(compile_exec( @@ -10929,6 +11662,127 @@ x = not True )); } + #[test] + fn test_plain_constant_bool_op_folds_to_selected_operand() { + let code = compile_exec( + "\ +x = 1 or 2 or 3 +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + let folded_small_int = code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::LoadSmallInt { i } + if i.get(OpArg::new(u32::from(u8::from(unit.arg)))) == 1 + ) + }); + let folded_const_one = code + .instructions + .iter() + .find_map(|unit| match unit.op { + Instruction::LoadConst { .. } => code.constants.get(usize::from(u8::from(unit.arg))), + _ => None, + }) + .is_some_and(|constant| { + matches!(constant, ConstantData::Integer { value } if *value == BigInt::from(1)) + }); + + assert!( + folded_small_int || folded_const_one, + "expected folded constant 1, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::Copy { .. } + | Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "plain constant BoolOp should not leave short-circuit scaffolding, got ops={ops:?}" + ); + } + + #[test] + fn test_starred_call_preserves_bool_op_short_circuit_shape() { + let code = compile_exec( + "\ +def f(g): + return g(*(() or (1,))) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.iter().any(|op| matches!(op, Instruction::Copy { .. })), + "starred BoolOp should keep short-circuit COPY, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!(op, Instruction::ToBool)), + "starred BoolOp should keep TO_BOOL, got ops={ops:?}" + ); + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::PopJumpIfTrue { .. })), + "starred BoolOp should keep POP_JUMP_IF_TRUE, got ops={ops:?}" + ); + } + + #[test] + fn test_partial_constant_bool_op_folds_prefix_in_value_context() { + let code = compile_exec( + "\ +def outer(null): + @False or null + def f(x): + pass +", + ); + let outer = find_code(&code, "outer").expect("missing outer code"); + let ops: Vec<_> = outer + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.iter().any(|op| { + matches!( + op, + Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. } + ) + }), + "expected surviving decorator expression to load null directly, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::Copy { .. } + | Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "partial constant BoolOp should not leave short-circuit scaffolding, got ops={ops:?}" + ); + } + #[test] fn test_nested_double_async_with() { assert_dis_snapshot!(compile_exec( @@ -11060,7 +11914,7 @@ def f(xs): } #[test] - fn test_builtin_tuple_list_set_genexpr_calls_are_optimized() { + fn test_builtin_tuple_genexpr_call_is_optimized_but_list_set_are_not() { let code = compile_exec( "\ def tuple_f(xs): @@ -11091,27 +11945,29 @@ def set_f(xs): assert_eq!(tuple_list_append, 2); let list_f = find_code(&code, "list_f").expect("missing list_f code"); - assert!(has_common_constant( - list_f, - bytecode::CommonConstant::BuiltinList - )); assert!( list_f .instructions .iter() - .any(|unit| matches!(unit.op, Instruction::ListAppend { .. })) + .any(|unit| matches!(unit.op, Instruction::Call { .. })), + "list(genexpr) should stay on the normal call path" + ); + assert!( + !has_common_constant(list_f, bytecode::CommonConstant::BuiltinList), + "CPython 3.14.2 does not optimize list(genexpr)" ); let set_f = find_code(&code, "set_f").expect("missing set_f code"); - assert!(has_common_constant( - set_f, - bytecode::CommonConstant::BuiltinSet - )); assert!( set_f .instructions .iter() - .any(|unit| matches!(unit.op, Instruction::SetAdd { .. })) + .any(|unit| matches!(unit.op, Instruction::Call { .. })), + "set(genexpr) should stay on the normal call path" + ); + assert!( + !has_common_constant(set_f, bytecode::CommonConstant::BuiltinSet), + "CPython 3.14.2 does not optimize set(genexpr)" ); } @@ -11401,9 +12257,9 @@ def f(node): .iter() .filter(|op| matches!(op, Instruction::ReturnValue)) .count(); - assert!( - return_count >= 3, - "expected multiple explicit return sites for shared final return case, got ops={ops:?}" + assert_eq!( + return_count, 5, + "expected cloned return sites for each shared return edge, got ops={ops:?}" ); } @@ -11458,31 +12314,776 @@ def f(parts): } #[test] - fn test_assert_without_message_raises_class_directly() { + fn test_for_exit_before_elif_does_not_leave_line_anchor_nop() { let code = compile_exec( "\ -def f(x): - assert x +from sys import maxsize +if maxsize == 2147483647: + for s in ('2147483648', '0o40000000000', '0x100000000', '0b10000000000000000000000000000000'): + try: + x = eval(s) + except OverflowError: + fail('OverflowError on huge integer literal %r' % s) +elif maxsize == 9223372036854775807: + pass ", ); - let f = find_code(&code, "f").expect("missing function code"); - let call_count = f + let ops: Vec<_> = code .instructions .iter() - .filter(|unit| matches!(unit.op, Instruction::Call { .. })) - .count(); - let push_null_count = f - .instructions - .iter() - .filter(|unit| matches!(unit.op, Instruction::PushNull)) - .count(); - - assert_eq!(call_count, 0); - assert_eq!(push_null_count, 0); - } + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); - #[test] - fn test_chained_compare_jump_uses_single_cleanup_copy() { + assert!( + ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }), + "expected for-exit epilogue without extra NOP, got ops={ops:?}" + ); + assert!( + !ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::Nop, + Instruction::LoadConst { .. }, + ] + ) + }), + "unexpected line-anchor NOP before for-exit epilogue, got ops={ops:?}" + ); + } + + #[test] + fn test_for_tuple_target_does_not_leave_loop_header_nop() { + let code = compile_exec( + "\ +def f(pairs): + for left, right in pairs: + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [ + Instruction::ForIter { .. }, + Instruction::UnpackSequence { .. } + ] + ) + }), + "expected FOR_ITER to flow directly into UNPACK_SEQUENCE, got ops={ops:?}" + ); + assert!( + !ops.windows(3).any(|window| { + matches!( + window, + [ + Instruction::ForIter { .. }, + Instruction::Nop, + Instruction::UnpackSequence { .. }, + ] + ) + }), + "unexpected loop-header NOP before tuple unpack, got ops={ops:?}" + ); + } + + #[test] + fn test_tstring_build_template_matches_cpython_stack_order() { + let code = compile_exec("t = t\"{0}\""); + let units: Vec<_> = code + .instructions + .iter() + .copied() + .filter(|unit| !matches!(unit.op, Instruction::Cache)) + .collect(); + + assert!( + units.windows(6).any(|window| { + matches!( + window, + [ + a, + b, + c, + d, + e, + f, + ] + if matches!(a.op, Instruction::LoadConst { .. }) + && matches!(b.op, Instruction::LoadSmallInt { .. }) + && matches!(c.op, Instruction::LoadConst { .. }) + && matches!(d.op, Instruction::BuildInterpolation { .. }) + && u8::from(d.arg) == 2 + && matches!(e.op, Instruction::BuildTuple { .. }) + && u8::from(e.arg) == 1 + && matches!(f.op, Instruction::BuildTemplate) + ) + }), + "expected CPython-style t-string lowering, got units={units:?}" + ); + assert!( + !units + .iter() + .any(|unit| matches!(unit.op, Instruction::Swap { .. })), + "unexpected SWAP in t-string lowering, got units={units:?}" + ); + } + + #[test] + fn test_tstring_debug_specifier_uses_debug_literal_and_repr_default() { + let code = compile_exec( + "\ +value = 42 +t = t\"Value: {value=}\" +", + ); + + let string_consts = code + .instructions + .iter() + .filter_map(|unit| match unit.op { + Instruction::LoadConst { consti } => { + Some(&code.constants[consti.get(OpArg::new(u32::from(u8::from(unit.arg))))]) + } + _ => None, + }) + .collect::>(); + + assert!( + string_consts.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if matches!( + &elements[..], + [ + ConstantData::Str { value: first }, + ConstantData::Str { value: second }, + ] if first.to_string() == "Value: value=" && second.is_empty() + ) + )), + "expected debug literal prefix in t-string constants, got {string_consts:?}" + ); + assert!( + code.instructions.iter().any(|unit| matches!( + unit.op, + Instruction::BuildInterpolation { .. } + ) && u8::from(unit.arg) == 10), + "expected default repr conversion for debug t-string" + ); + } + + #[test] + fn test_tstring_literal_preserves_surrogate_wtf8() { + let code = compile_exec("t = t\"\\ud800\""); + + assert!(code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Str { value } if value.clone().into_bytes() == [0xED, 0xA0, 0x80] + ))); + } + + #[test] + fn test_break_in_finally_after_return_keeps_load_fast_check_for_loop_locals() { + let code = compile_exec( + "\ +def g2(x): + for count in [0, 1]: + for count2 in [10, 20]: + try: + return count + count2 + finally: + if x: + break + return 'end', count, count2 +", + ); + let g2 = find_code(&code, "g2").expect("missing g2 code"); + let ops: Vec<_> = g2 + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::LoadConst { .. }, + Instruction::LoadFastCheck { .. }, + Instruction::LoadFastCheck { .. }, + Instruction::BuildTuple { .. }, + ] + ) + }), + "expected LOAD_FAST_CHECK pair for after-return loop locals, got ops={ops:?}" + ); + } + + #[test] + fn test_assert_without_message_raises_class_directly() { + let code = compile_exec( + "\ +def f(x): + assert x +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let call_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::Call { .. })) + .count(); + let push_null_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::PushNull)) + .count(); + + assert_eq!(call_count, 0); + assert_eq!(push_null_count, 0); + } + + #[test] + fn test_assert_with_message_uses_common_constant_direct_call() { + let code = compile_exec( + "\ +def f(x, y): + assert x, y +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let load_assertion = f + .instructions + .iter() + .position(|unit| { + matches!(unit.op, Instruction::LoadCommonConstant { .. }) + && matches!( + unit.op, + Instruction::LoadCommonConstant { idx } + if idx.get(OpArg::new(u32::from(u8::from(unit.arg)))) + == bytecode::CommonConstant::AssertionError + ) + }) + .expect("missing LOAD_COMMON_CONSTANT AssertionError"); + + assert!( + !matches!( + f.instructions.get(load_assertion + 1).map(|unit| unit.op), + Some(Instruction::PushNull) + ), + "assert message path should not use PUSH_NULL, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + matches!( + f.instructions.get(load_assertion + 2).map(|unit| unit.op), + Some(Instruction::Call { .. }) + ), + "expected direct CALL after loading assert message, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + + let call_arg = f.instructions[load_assertion + 2].arg; + assert_eq!(u8::from(call_arg), 0); + } + + #[test] + fn test_bare_function_annotations_check_attribute_and_subscript_expressions() { + assert_dis_snapshot!(compile_exec( + "\ +def f(one: int): + int.new_attr: int + [list][0].new_attr: [int, str] + my_lst = [1] + my_lst[one]: int + return my_lst +" + )); + } + + #[test] + fn test_non_simple_bare_name_annotation_does_not_create_local_binding() { + let code = compile_exec( + "\ +def f2bad(): + (no_such_global): int + print(no_such_global) +", + ); + let f = find_code(&code, "f2bad").expect("missing f2bad code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadGlobal { .. })), + "expected LOAD_GLOBAL for non-simple bare annotated name, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + !f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadFastCheck { .. })), + "non-simple bare annotated name should not become a local binding, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_true_if_pass_keeps_line_anchor_nop() { + assert_dis_snapshot!(compile_exec( + "\ +if 1: + pass +" + )); + } + + #[test] + fn test_negative_constant_binop_folds_after_unary_folding() { + let code = compile_exec( + "\ +def f(): + return -2147483647 - 1 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + assert!( + !f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BinaryOp { .. })), + "negative constant expression should fold to a single constant, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadConst { .. })), + "expected folded constant load, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_genexpr_filter_header_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(it): + return (x for x in it if x) +", + ); + let genexpr = find_code(&code, "").expect("missing code"); + let store_fast_load_fast_idx = genexpr + .instructions + .iter() + .position(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })) + .expect("missing STORE_FAST_LOAD_FAST in genexpr header"); + + assert!( + matches!( + genexpr + .instructions + .get(store_fast_load_fast_idx + 1) + .map(|unit| unit.op), + Some(Instruction::ToBool) + ), + "expected TO_BOOL immediately after STORE_FAST_LOAD_FAST, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_multi_with_header_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(manager): + with manager() as x, manager(): + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in multi-with header, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_sequential_store_then_load_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(self): + x = ''; y = \"\"; self.assertTrue(len(x) == 0 and x == y) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in sequential statement body, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_match_guard_capture_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(): + match 0: + case x if x: + z = 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in match guard capture path, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_match_nested_capture_uses_store_fast_store_fast() { + let code = compile_exec( + "\ +def f(x): + match x: + case ((0 as w) as z): + return w, z +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastStoreFast { .. })), + "expected STORE_FAST_STORE_FAST in nested match capture path, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_match_value_real_zero_minus_zero_complex_folds_to_negative_zero_imag() { + let code = compile_exec( + "\ +def f(x): + match x: + case 0 - 0j: + return 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.constants.iter().any(|constant| matches!( + constant, + ConstantData::Complex { value } + if value.re == 0.0 && value.im == 0.0 && value.im.is_sign_negative() + )), + "expected folded -0j constant in match value" + ); + } + + #[test] + fn test_match_or_uses_shared_success_block() { + let code = compile_exec( + "\ +def http_error(status): + match status: + case 400: + return 'Bad request' + case 401 | 403 | 404: + return 'Not allowed' + case 418: + return 'I am a teapot' +", + ); + let f = find_code(&code, "http_error").expect("missing http_error code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let jump_positions: Vec<_> = ops + .iter() + .enumerate() + .filter_map(|(i, op)| matches!(op, Instruction::JumpForward { .. }).then_some(i)) + .collect(); + + assert!( + jump_positions.len() >= 4, + "expected shared-success JumpForward ops in OR pattern, got ops={ops:?}" + ); + + let first_pop_top_pair = ops + .windows(2) + .position(|window| matches!(window, [Instruction::PopTop, Instruction::PopTop])) + .expect("missing POP_TOP/POP_TOP success cleanup"); + + assert!( + jump_positions + .iter() + .take(3) + .all(|&idx| idx < first_pop_top_pair), + "expected OR-alternative jumps before shared success cleanup, got ops={ops:?}" + ); + } + + #[test] + fn test_match_mapping_attribute_key_keeps_plain_load_fast() { + let code = compile_exec( + "\ +def f(self): + class Keys: + KEY = 'a' + x = {'a': 0, 'b': 1} + with self.assertRaises(ValueError): + match x: + case {Keys.KEY: y, 'a': z}: + w = 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let key_load_idx = f + .instructions + .iter() + .position(|unit| match unit.op { + Instruction::LoadAttr { namei } => { + let load_attr = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + f.names[usize::try_from(load_attr.name_idx()).unwrap()].as_str() == "KEY" + } + _ => false, + }) + .expect("missing Keys.KEY attribute load"); + let prev = f.instructions[key_load_idx - 1].op; + assert!( + matches!(prev, Instruction::LoadFast { .. }), + "expected plain LOAD_FAST before Keys.KEY mapping key, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + #[ignore = "debug trace for sequence star-wildcard pattern layout"] + fn test_debug_trace_match_sequence_star_wildcard_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(w): + match w: + case [x, *_, y]: + z = 0 + return x, y, z +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_genexpr_true_filter_omits_bool_scaffolding() { + let code = compile_exec( + "\ +def f(it): + return (x for x in it if True) +", + ); + let genexpr = find_code(&code, "").expect("missing code"); + assert!( + !genexpr.instructions.iter().any(|unit| { + matches!(unit.op, Instruction::LoadConst { .. }) + && matches!( + genexpr.constants.get(usize::from(u8::from(unit.arg))), + Some(ConstantData::Boolean { value: true }) + ) + }), + "constant-true filter should not load True, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + !genexpr + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::PopJumpIfTrue { .. })), + "constant-true filter should not leave POP_JUMP_IF_TRUE scaffolding, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_classdictcell_uses_load_closure_path_and_borrows_after_optimize() { + let code = compile_exec( + "\ +class C: + def method(self): + return 1 +", + ); + let class_code = find_code(&code, "C").expect("missing class code"); + let store_classdictcell = class_code + .instructions + .iter() + .position(|unit| { + matches!( + unit.op, + Instruction::StoreName { namei } + if class_code.names + [namei.get(OpArg::new(u32::from(u8::from(unit.arg)))) as usize] + .as_str() + == "__classdictcell__" + ) + }) + .expect("missing STORE_NAME __classdictcell__"); + + assert!( + matches!( + class_code + .instructions + .get(store_classdictcell.saturating_sub(1)) + .map(|unit| unit.op), + Some(Instruction::LoadFastBorrow { .. }) + ), + "expected LOAD_FAST_BORROW before __classdictcell__ store, got ops={:?}", + class_code + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_future_annotations_class_keeps_conditional_annotations_cell() { + let code = compile_exec( + "\ +from __future__ import annotations +class C: + x: int +", + ); + let class_code = find_code(&code, "C").expect("missing class code"); + + assert!( + class_code + .cellvars + .iter() + .any(|name| name.as_str() == "__conditional_annotations__"), + "expected __conditional_annotations__ cellvar, got cellvars={:?}", + class_code.cellvars + ); + } + + #[test] + fn test_plain_super_call_keeps_class_freevar() { + let code = compile_exec( + "\ +class A: + pass + +class B(A): + def method(self): + return super() +", + ); + let method = find_code(&code, "method").expect("missing method code"); + assert!( + method.freevars.iter().any(|name| name == "__class__"), + "plain super() must keep __class__ freevar, got freevars={:?}", + method.freevars + ); + assert!( + method + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::CopyFreeVars { .. })), + "plain super() must keep COPY_FREE_VARS prelude, got ops={:?}", + method + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_chained_compare_jump_uses_single_cleanup_copy() { let code = compile_exec( "\ def f(code): @@ -11583,6 +13184,120 @@ def f(): ); } + #[test] + fn test_named_except_cleanup_keeps_jump_over_cleanup_and_next_try() { + let code = compile_exec( + r#" +def f(self): + try: + assert 0, 'msg' + except AssertionError as e: + self.assertEqual(e.args[0], 'msg') + else: + self.fail("AssertionError not raised by assert 0") + + try: + assert False + except AssertionError as e: + self.assertEqual(len(e.args), 0) + else: + self.fail("AssertionError not raised by 'assert False'") +"#, + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let first_pop_except = ops + .iter() + .position(|op| matches!(op, Instruction::PopExcept)) + .expect("missing POP_EXCEPT"); + let window = &ops[first_pop_except..(first_pop_except + 6).min(ops.len())]; + assert!( + matches!( + window, + [ + Instruction::PopExcept, + Instruction::LoadConst { .. }, + Instruction::StoreName { .. } | Instruction::StoreFast { .. }, + Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }, + Instruction::JumpForward { .. }, + .. + ] + ), + "expected named except cleanup to jump over cleanup reraise block, got ops={window:?}" + ); + } + + #[test] + fn test_bare_except_deopts_post_handler_load_fast_borrow() { + let code = compile_exec( + "\ +def f(self): + try: + 1 / 0 + except: + pass + with self.assertRaises(SyntaxError): + pass +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let attr_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing LOAD_ATTR for assertRaises"); + assert!( + matches!(ops.get(attr_idx - 1), Some(Instruction::LoadFast { .. })), + "bare except tail should deopt self to LOAD_FAST, got ops={ops:?}" + ); + } + + #[test] + fn test_typed_except_keeps_post_handler_load_fast_borrow() { + let code = compile_exec( + "\ +def f(self): + try: + 1 / 0 + except ZeroDivisionError: + pass + with self.assertRaises(SyntaxError): + pass +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let attr_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing LOAD_ATTR for assertRaises"); + assert!( + matches!( + ops.get(attr_idx - 1), + Some(Instruction::LoadFastBorrow { .. }) + ), + "typed except tail should keep LOAD_FAST_BORROW, got ops={ops:?}" + ); + } + #[test] fn test_constant_slice_folding_handles_string_and_bigint_bounds() { let code = compile_exec( @@ -11625,7 +13340,201 @@ def f(names, cls): .filter(|unit| matches!(unit.op, Instruction::ReturnValue)) .count(); - assert_eq!(return_count, 1); + assert_eq!(return_count, 1); + } + + #[test] + fn test_non_none_final_return_is_not_duplicated() { + let code = compile_exec( + "\ +def f(p, s): + if p == '': + if s == '': + return 0 + return -1 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let minus_one_loads = f + .instructions + .iter() + .filter(|unit| { + matches!( + unit.op, + Instruction::LoadConst { consti } + if matches!( + f.constants.get( + consti + .get(OpArg::new(u32::from(u8::from(unit.arg)))) + .as_usize() + ), + Some(ConstantData::Integer { value }) if value == &BigInt::from(-1) + ) + ) + }) + .count(); + + assert_eq!( + minus_one_loads, + 1, + "expected a single final return -1 epilogue, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_named_except_conditional_branch_duplicates_cleanup_return() { + let code = compile_exec( + "\ +def f(self): + try: + raise TypeError('x') + except TypeError as e: + if '+' not in str(e): + self.fail('join() ate exception message') +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let cleanup_return_count = ops + .windows(6) + .filter(|window| { + matches!( + window, + [ + Instruction::PopExcept, + Instruction::LoadConst { .. }, + Instruction::StoreFast { .. } | Instruction::StoreName { .. }, + Instruction::DeleteFast { .. } | Instruction::DeleteName { .. }, + Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }) + .count(); + + assert_eq!( + cleanup_return_count, 2, + "expected duplicated named-except cleanup return blocks, got ops={ops:?}" + ); + } + + #[test] + fn test_listcomp_cleanup_tail_keeps_split_store_fast_pair() { + let code = compile_exec( + "\ +def f(escaped_string, quote_types): + possible_quotes = [q for q in quote_types if q not in escaped_string] + return possible_quotes +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let pop_iter_idx = ops + .iter() + .position(|op| matches!(op, Instruction::PopIter)) + .expect("missing POP_ITER"); + let tail = &ops[pop_iter_idx + 1..]; + + assert!( + matches!( + tail, + [ + Instruction::StoreFast { .. }, + Instruction::StoreFast { .. }, + Instruction::LoadFastBorrow { .. }, + Instruction::ReturnValue, + .. + ] + ), + "expected split STORE_FAST pair after listcomp cleanup, got ops={ops:?}" + ); + } + + #[test] + fn test_with_suppress_tail_duplicates_final_return_none() { + let code = compile_exec( + "\ +def f(cm, cond): + if cond: + with cm(): + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let return_count = ops + .iter() + .filter(|op| matches!(op, Instruction::ReturnValue)) + .count(); + + assert_eq!( + return_count, 3, + "expected duplicated return-none epilogues, got ops={ops:?}" + ); + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::JumpBackwardNoInterrupt { .. })), + "with suppress tail should not jump back to shared return block, got ops={ops:?}" + ); + } + + #[test] + fn test_genexpr_compare_header_keeps_split_store_then_borrow_load() { + let code = compile_exec( + "\ +def f(it): + return (offset == (4, 10) for offset in it) +", + ); + let genexpr = find_code(&code, "").expect("missing code"); + let ops: Vec<_> = genexpr + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::StoreFastLoadFast { .. })), + "expected compare header to keep split STORE_FAST/LOAD_FAST_BORROW, got ops={ops:?}" + ); + assert!( + ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::StoreFast { .. }, + Instruction::LoadFastBorrow { .. }, + Instruction::LoadConst { .. }, + Instruction::CompareOp { .. }, + ] + ) + }), + "expected split compare header sequence, got ops={ops:?}" + ); } #[test] @@ -11745,6 +13654,74 @@ def f(x): })); } + #[test] + fn test_string_and_bytes_binops_constant_fold_like_cpython() { + let code = compile_exec( + "\ +x = b'\\\\' + b'u1881'\n\ +y = 103 * 'a' + 'x'\n", + ); + + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BinaryOp { .. })), + "unexpected runtime BINARY_OP in folded string/bytes constants: {:?}", + code.instructions + ); + assert!(code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Bytes { value } if value == b"\\u1881" + ))); + let expected = format!("{}x", "a".repeat(103)); + assert!(code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Str { value } + if value.to_string() == expected + ))); + } + + #[test] + fn test_constant_string_subscript_folds_inside_collection() { + let code = compile_exec( + "\ +values = [item for item in [r\"\\\\'a\\\\'\", r\"\\t3\", r\"\\\\\"[0]]]\n", + ); + + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BinaryOp { .. })), + "unexpected runtime BINARY_OP after constant subscript folding: {:?}", + code.instructions + ); + assert!(code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if elements.len() == 3 + && matches!(&elements[2], ConstantData::Str { value } if value.to_string() == "\\") + ))); + } + + #[test] + fn test_constant_string_subscript_with_surrogate_skips_lossy_fold() { + let code = compile_exec("value = \"\\ud800\"[0]\n"); + + assert!( + code.instructions.iter().any(|unit| match unit.op { + Instruction::BinaryOp { op } => { + op.get(OpArg::new(u32::from(u8::from(unit.arg)))) + == oparg::BinaryOperator::Subscr + } + _ => false, + }), + "expected runtime subscript for surrogate literal, got instructions={:?}", + code.instructions + ); + } + #[test] fn test_list_of_constant_tuples_uses_list_extend() { let code = compile_exec( @@ -11812,7 +13789,7 @@ def f(): ); assert!(f.constants.iter().any(|constant| matches!( constant, - ConstantData::Tuple { elements } + ConstantData::Frozenset { elements } if matches!( elements.as_slice(), [ @@ -11824,6 +13801,306 @@ def f(): ))); } + #[test] + fn test_starred_tuple_iterable_drops_list_to_tuple_before_get_iter() { + let code = compile_exec( + "\ +def f(a, b, c): + for x in *a, *b, *c: + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + assert!( + !has_intrinsic_1(f, IntrinsicFunction1::ListToTuple), + "LIST_TO_TUPLE should be removed before GET_ITER in for-iterable context" + ); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::GetIter)), + "expected GET_ITER in for loop" + ); + } + + #[test] + fn test_comprehension_single_list_iterable_uses_tuple() { + let code = compile_exec( + "\ +def g(): + [x for x in [(yield 1)]] +", + ); + let g = find_code(&code, "g").expect("missing g code"); + let ops: Vec<_> = g + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [Instruction::BuildTuple { .. }, Instruction::GetIter] + ) + }), + "expected BUILD_TUPLE before GET_ITER for single-item list iterable in comprehension, got ops={ops:?}" + ); + } + + #[test] + fn test_nested_comprehension_list_iterable_uses_tuple() { + let code = compile_exec( + "\ +def f(): + return [[y for y in [x, x + 1]] for x in [1, 3, 5]] +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [Instruction::BuildTuple { .. }, Instruction::GetIter] + ) + }), + "expected BUILD_TUPLE before GET_ITER for nested list iterable in comprehension, got ops={ops:?}" + ); + } + + #[test] + fn test_constant_comprehension_iterable_with_unary_int_uses_tuple_const() { + let code = compile_exec( + "\ +l = lambda : [2 < x for x in [-1, 3, 0]] +", + ); + let lambda = find_code(&code, "").expect("missing lambda code"); + + assert!( + lambda.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if matches!( + elements.as_slice(), + [ + ConstantData::Integer { .. }, + ConstantData::Integer { .. }, + ConstantData::Integer { .. } + ] + ) + )), + "expected folded tuple constant for comprehension iterable" + ); + } + + #[test] + fn test_constant_false_while_else_deopts_post_else_borrows() { + let code = compile_exec( + "\ +def f(self): + x = 0 + while 0: + x = 1 + else: + x = 2 + self.assertEqual(x, 2) +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + let assert_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing assertEqual call"); + let window = &ops[assert_idx.saturating_sub(1)..(assert_idx + 3).min(ops.len())]; + assert!( + matches!( + window, + [ + Instruction::LoadFast { .. }, + Instruction::LoadAttr { .. }, + Instruction::LoadFast { .. }, + .. + ] + ), + "expected post-else assertEqual call to use plain LOAD_FAST, got ops={window:?}" + ); + } + + #[test] + fn test_single_unpack_assignment_disables_constant_collection_folding() { + let code = compile_exec("a, b, c = 1, 2, 3\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!(unit.op, Instruction::UnpackSequence { .. }) + || matches!(unit.op, Instruction::LoadConst { .. }) + && matches!( + code.constants.get(usize::from(u8::from(unit.arg))), + Some(ConstantData::Tuple { .. }) + ) + }), + "single unpack assignment should keep builder form for later lowering, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::LoadSmallInt { .. })) + .count() + >= 3, + "expected individual constant loads before unpack-target stores, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_chained_unpack_assignment_keeps_constant_collection_folding() { + let code = compile_exec("(a, b) = c = d = (1, 2)\n"); + + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadConst { .. })), + "chained unpack assignment should keep tuple constant, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::UnpackSequence { .. })), + "chained unpack assignment should still unpack the copied tuple, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_true_assert_skips_message_nested_scope() { + let code = compile_exec("assert 1, (lambda x: x + 1)\n"); + + assert_eq!( + code.constants + .iter() + .filter(|constant| matches!(constant, ConstantData::Code { .. })) + .count(), + 0, + "constant-true assert should not compile the skipped message lambda" + ); + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::RaiseVarargs { .. })), + "constant-true assert should be elided, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_false_assert_uses_direct_raise_shape() { + let code = compile_exec("assert 0, (lambda x: x + 1)\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "constant-false assert should use direct raise shape, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::RaiseVarargs { .. })), + "constant-false assert should still raise, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert_eq!( + code.constants + .iter() + .filter(|constant| matches!(constant, ConstantData::Code { .. })) + .count(), + 1, + "constant-false assert should still compile the message lambda" + ); + } + + #[test] + fn test_constant_unary_positive_and_invert_fold() { + let code = compile_exec("x = +1\nx = ~1\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::CallIntrinsic1 { .. } | Instruction::UnaryInvert + ) + }), + "constant unary ops should fold away, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_bool_invert_is_not_const_folded() { + let code = compile_exec("x = ~True\n"); + + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::UnaryInvert)), + "~bool should remain unfurled to match CPython, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + #[test] fn test_optimized_assert_preserves_nested_scope_order() { compile_exec_optimized( @@ -11864,4 +14141,40 @@ def f(items): ", ); } + + #[test] + fn test_try_else_nested_scopes_keep_subtable_cursor_aligned() { + let code = compile_exec( + "\ +try: + import missing_mod +except ImportError: + def fallback(): + return 0 +else: + def impl(): + return reversed('abc') +", + ); + + assert!( + find_code(&code, "fallback").is_some(), + "missing fallback code" + ); + let impl_code = find_code(&code, "impl").expect("missing impl code"); + assert!( + impl_code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::LoadGlobal { .. } | Instruction::LoadName { .. } + ) + }), + "expected impl to compile global name access, got ops={:?}", + impl_code + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 61e549199d5..7d09c2cef83 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -3,15 +3,17 @@ use core::ops; use crate::{IndexMap, IndexSet, error::InternalError}; use malachite_bigint::BigInt; +use num_complex::Complex; use num_traits::{ToPrimitive, Zero}; use rustpython_compiler_core::{ OneIndexed, SourceLocation, bytecode::{ - AnyInstruction, AnyOpcode, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, + AnyInstruction, AnyOpcode, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, - InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg, Opcode, - PseudoInstruction, PseudoOpcode, PyCodeLocationInfoKind, encode_exception_table, oparg, + InstrDisplayContext, Instruction, InstructionMetadata, IntrinsicFunction1, Label, OpArg, + Opcode, PseudoInstruction, PseudoOpcode, PyCodeLocationInfoKind, encode_exception_table, + oparg, }, varint::{write_signed_varint, write_varint}, }; @@ -108,6 +110,7 @@ pub struct InstructionInfo { pub location: SourceLocation, pub end_location: SourceLocation, pub except_handler: Option, + pub folded_from_nonliteral_expr: bool, /// Override line number for linetable (e.g., line 0 for module RESUME) pub lineno_override: Option, /// Number of CACHE code units emitted after this instruction @@ -129,6 +132,7 @@ fn set_to_nop(info: &mut InstructionInfo) { info.instr = Instruction::Nop.into(); info.arg = OpArg::new(0); info.target = BlockIdx::NULL; + info.folded_from_nonliteral_expr = false; info.cache_entries = 0; } @@ -148,6 +152,8 @@ pub struct Block { pub start_depth: Option, /// Whether this block is only reachable via exception table (b_cold) pub cold: bool, + /// Whether LOAD_FAST borrow optimization should be suppressed for this block. + pub disable_load_fast_borrow: bool, } impl Default for Block { @@ -159,6 +165,7 @@ impl Default for Block { preserve_lasti: false, start_depth: None, cold: false, + disable_load_fast_borrow: false, } } } @@ -170,6 +177,7 @@ pub struct CodeInfo { pub blocks: Vec, pub current_block: BlockIdx, + pub annotations_blocks: Option>, pub metadata: CodeUnitMetadata, @@ -199,20 +207,17 @@ impl CodeInfo { mut self, opts: &crate::compile::CompileOpts, ) -> crate::InternalResult { + self.splice_annotations_blocks(); // Constant folding passes self.fold_binop_constants(); - self.remove_nops(); - self.fold_unary_negative(); + self.fold_unary_constants(); self.fold_binop_constants(); // re-run after unary folding: -1 + 2 → 1 - self.remove_nops(); // remove NOPs so tuple/list/set see contiguous LOADs self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); - self.remove_nops(); // remove NOPs from collection folding self.fold_const_iterable_for_iter(); self.convert_to_load_small_int(); self.remove_unused_consts(); - self.remove_nops(); // DCE always runs (removes dead code after terminal instructions) self.dce(); @@ -223,7 +228,7 @@ impl CodeInfo { self.eliminate_dead_stores(); // apply_static_swaps: reorder stores to eliminate SWAPs self.apply_static_swaps(); - // Peephole optimizer creates superinstructions matching CPython + // Peephole optimizer handles constant and compare folding. self.peephole_optimize(); // Phase 1: _PyCfg_OptimizeCodeUnit (flowgraph.c) @@ -235,7 +240,11 @@ impl CodeInfo { jump_threading(&mut self.blocks); self.eliminate_unreachable_blocks(); self.remove_nops(); - // TODO: insert_superinstructions disabled pending StoreFastLoadFast VM fix + self.add_checks_for_loads_of_uninitialized_variables(); + // CPython inserts superinstructions in _PyCfg_OptimizeCodeUnit, before + // later jump normalization / block reordering can create adjacencies + // that never exist at this stage in flowgraph.c. + self.insert_superinstructions(); push_cold_blocks_to_end(&mut self.blocks); // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) @@ -247,9 +256,11 @@ impl CodeInfo { self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions self.eliminate_unreachable_blocks(); resolve_line_numbers(&mut self.blocks); - duplicate_end_returns(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); + duplicate_end_returns(&mut self.blocks, &self.metadata); self.dce(); // truncate after terminal in blocks that got return duplicated self.eliminate_unreachable_blocks(); // remove now-unreachable last block + self.remove_redundant_const_pop_top_pairs(); remove_redundant_nops_and_jumps(&mut self.blocks); // Some jump-only blocks only appear after late CFG cleanup. Thread them // once more so loop backedges stay direct instead of becoming @@ -260,12 +271,44 @@ impl CodeInfo { reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); - self.add_checks_for_loads_of_uninitialized_variables(); + inline_with_suppress_return_blocks(&mut self.blocks); + inline_pop_except_return_blocks(&mut self.blocks); + duplicate_named_except_cleanup_returns(&mut self.blocks, &self.metadata); + self.eliminate_unreachable_blocks(); + // Late CFG cleanup can create new same-line STORE_FAST/LOAD_FAST and + // STORE_FAST/STORE_FAST adjacencies in match/capture code paths that + // did not exist during the earlier flowgraph-like pass. + self.insert_superinstructions(); + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + // Late CFG cleanup can create or reshuffle handler entry blocks. + // Refresh exceptional block flags before optimize_load_fast_borrow so + // borrow loads are not introduced into exception-handler paths. + mark_except_handlers(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); + // CPython's optimize_load_fast runs with block start depths already known. + // Compute them here so the abstract stack simulation can use the real + // CFG entry depth for each block. + let max_stackdepth = self.max_stackdepth()?; + // Match CPython order: pseudo ops are lowered after stackdepth + // calculation but before optimize_load_fast. + convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); + self.compute_load_fast_start_depths(); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); + self.deoptimize_borrow_after_push_exc_info(); + self.deoptimize_borrow_for_handler_return_paths(); + self.deoptimize_borrow_for_match_keys_attr(); + self.deoptimize_store_fast_store_fast_after_cleanup(); + self.apply_static_swaps(); + self.insert_superinstructions(); + self.deoptimize_store_fast_store_fast_after_cleanup(); self.optimize_load_global_push_null(); - - let max_stackdepth = self.max_stackdepth()?; + self.reorder_entry_prefix_cell_setup(); + self.remove_unused_consts(); let Self { flags, @@ -274,6 +317,7 @@ impl CodeInfo { mut blocks, current_block: _, + annotations_blocks: _, metadata, static_attributes: _, in_inlined_comp: _, @@ -308,6 +352,7 @@ impl CodeInfo { // Convert pseudo ops (LoadClosure uses cellfixedoffsets) and fixup DEREF opargs convert_pseudo_ops(&mut blocks, &cellfixedoffsets); fixup_deref_opargs(&mut blocks, &cellfixedoffsets); + deoptimize_borrow_after_push_exc_info_in_blocks(&mut blocks); // Remove redundant NOPs, keeping line-marker NOPs only when // they are needed to preserve tracing. let mut block_order = Vec::new(); @@ -330,28 +375,25 @@ impl CodeInfo { let mut remove = false; if matches!(instr.instr.real(), Some(Instruction::Nop)) { - // Remove location-less NOPs. if lineno < 0 || prev_lineno == lineno { remove = true; - } - // Remove if the next instruction has same line or no line. - else if src < src_instructions.len() - 1 { - let next_lineno = - src_instructions[src + 1] + } else if src < src_instructions.len() - 1 { + if src_instructions[src + 1].folded_from_nonliteral_expr { + remove = true; + } else { + let next_lineno = src_instructions[src + 1] .lineno_override .unwrap_or_else(|| { src_instructions[src + 1].location.line.get() as i32 }); - if next_lineno == lineno { - remove = true; - } else if next_lineno < 0 { - src_instructions[src + 1].lineno_override = Some(lineno); - remove = true; + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } } - } - // Last instruction in block: compare with first real location - // in the next non-empty block. - else { + } else { let mut next = blocks[bi].next; while next != BlockIdx::NULL && blocks[next.idx()].instructions.is_empty() { next = blocks[next.idx()].next; @@ -620,6 +662,63 @@ impl CodeInfo { } } + fn reorder_entry_prefix_cell_setup(&mut self) { + let Some(entry) = self.blocks.first_mut() else { + return; + }; + let ncells = self.metadata.cellvars.len(); + let nfrees = self.metadata.freevars.len(); + if ncells == 0 && nfrees == 0 { + return; + } + + let prefix_len = entry + .instructions + .iter() + .take_while(|info| { + matches!( + info.instr.real(), + Some(Instruction::MakeCell { .. } | Instruction::CopyFreeVars { .. }) + ) + }) + .count(); + if prefix_len == 0 { + return; + } + + let original_prefix = entry.instructions[..prefix_len].to_vec(); + let anchor = original_prefix[0]; + let rest = entry.instructions.split_off(prefix_len); + entry.instructions.clear(); + + if nfrees > 0 { + entry.instructions.push(InstructionInfo { + instr: Instruction::CopyFreeVars { n: Arg::marker() }.into(), + arg: OpArg::new(nfrees as u32), + ..anchor + }); + } + + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + let mut sorted = vec![None; self.metadata.varnames.len() + ncells]; + for (oldindex, fixed) in cellfixedoffsets.iter().copied().take(ncells).enumerate() { + sorted[fixed as usize] = Some(oldindex); + } + for oldindex in sorted.into_iter().flatten() { + entry.instructions.push(InstructionInfo { + instr: Instruction::MakeCell { i: Arg::marker() }.into(), + arg: OpArg::new(oldindex as u32), + ..anchor + }); + } + + entry.instructions.extend(rest); + } + /// Clear blocks that are unreachable (not entry, not a jump target, /// and only reachable via fall-through from a terminal block). fn eliminate_unreachable_blocks(&mut self) { @@ -668,51 +767,103 @@ impl CodeInfo { } } - /// Fold LOAD_CONST/LOAD_SMALL_INT + UNARY_NEGATIVE → LOAD_CONST (negative value) - fn fold_unary_negative(&mut self) { + fn eval_unary_constant( + operand: &ConstantData, + op: Instruction, + intrinsic: Option, + ) -> Option { + match (operand, op, intrinsic) { + (ConstantData::Integer { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Integer { value: -value }) + } + (ConstantData::Float { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Float { value: -value }) + } + (ConstantData::Complex { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Complex { value: -value }) + } + (ConstantData::Boolean { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Integer { + value: BigInt::from(-i32::from(*value)), + }) + } + (ConstantData::Integer { value }, Instruction::UnaryInvert, None) => { + Some(ConstantData::Integer { value: !value }) + } + (ConstantData::Boolean { .. }, Instruction::UnaryInvert, None) => None, + ( + ConstantData::Integer { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Integer { + value: value.clone(), + }), + ( + ConstantData::Float { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Float { value: *value }), + ( + ConstantData::Boolean { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Integer { + value: BigInt::from(i32::from(*value)), + }), + _ => None, + } + } + + /// Fold constant unary operations following CPython fold_const_unaryop(). + fn fold_unary_constants(&mut self) { for block in &mut self.blocks { let mut i = 0; - while i + 1 < block.instructions.len() { - let next = &block.instructions[i + 1]; - let Some(Instruction::UnaryNegative) = next.instr.real() else { + while i < block.instructions.len() { + let instr = &block.instructions[i]; + let (op, intrinsic) = match instr.instr.real() { + Some(Instruction::UnaryNegative) => (Instruction::UnaryNegative, None), + Some(Instruction::UnaryInvert) => (Instruction::UnaryInvert, None), + Some(Instruction::CallIntrinsic1 { func }) + if matches!( + func.get(instr.arg), + oparg::IntrinsicFunction1::UnaryPositive + ) => + { + ( + Instruction::CallIntrinsic1 { + func: Arg::marker(), + }, + Some(func.get(instr.arg)), + ) + } + _ => { + i += 1; + continue; + } + }; + let Some(operand_index) = i + .checked_sub(1) + .and_then(|start| Self::get_const_loading_instr_indices(block, start, 1)) + .and_then(|indices| indices.into_iter().next()) + else { i += 1; continue; }; - let curr = &block.instructions[i]; - let value = match curr.instr.real() { - Some(Instruction::LoadConst { .. }) => { - let idx = u32::from(curr.arg) as usize; - match self.metadata.consts.get_index(idx) { - Some(ConstantData::Integer { value }) => { - Some(ConstantData::Integer { value: -value }) - } - Some(ConstantData::Float { value }) => { - Some(ConstantData::Float { value: -value }) - } - _ => None, - } - } - Some(Instruction::LoadSmallInt { .. }) => { - let v = u32::from(curr.arg) as i32; - Some(ConstantData::Integer { - value: BigInt::from(-v), - }) + let operand = + Self::get_const_value_from(&self.metadata, &block.instructions[operand_index]); + if let Some(operand) = operand + && let Some(folded_const) = Self::eval_unary_constant(&operand, op, intrinsic) + { + let (const_idx, _) = self.metadata.consts.insert_full(folded_const); + let folded_from_nonliteral_expr = true; + set_to_nop(&mut block.instructions[operand_index]); + block.instructions[i].instr = Instruction::LoadConst { + consti: Arg::marker(), } - _ => None, - }; - if let Some(neg_const) = value { - let (const_idx, _) = self.metadata.consts.insert_full(neg_const); - // Replace LOAD_CONST/LOAD_SMALL_INT with new LOAD_CONST - let load_location = block.instructions[i].location; - block.instructions[i].instr = Opcode::LoadConst.into(); + .into(); block.instructions[i].arg = OpArg::new(const_idx as u32); - // Replace UNARY_NEGATIVE with NOP, inheriting the LOAD_CONST - // location so that remove_nops can clean it up - set_to_nop(&mut block.instructions[i + 1]); - block.instructions[i + 1].location = load_location; - block.instructions[i + 1].end_location = block.instructions[i].end_location; - // Skip the NOP, don't re-check - i += 2; + block.instructions[i].folded_from_nonliteral_expr = folded_from_nonliteral_expr; + i = i.saturating_sub(1); } else { i += 1; } @@ -720,6 +871,27 @@ impl CodeInfo { } } + fn get_const_loading_instr_indices( + block: &Block, + mut start: usize, + size: usize, + ) -> Option> { + let mut indices = Vec::with_capacity(size); + loop { + let instr = block.instructions.get(start)?; + if !matches!(instr.instr.real(), Some(Instruction::Nop)) { + Self::get_const_value_from_dummy(instr)?; + indices.push(start); + if indices.len() == size { + break; + } + } + start = start.checked_sub(1)?; + } + indices.reverse(); + Some(indices) + } + /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + LOAD_CONST/LOAD_SMALL_INT + BINARY_OP /// into a single LOAD_CONST when the result is computable at compile time. /// = fold_binops_on_constants in CPython flowgraph.c @@ -728,22 +900,34 @@ impl CodeInfo { for block in &mut self.blocks { let mut i = 0; - while i + 2 < block.instructions.len() { - // Check pattern: LOAD_CONST/LOAD_SMALL_INT, LOAD_CONST/LOAD_SMALL_INT, BINARY_OP - let Some(Instruction::BinaryOp { .. }) = block.instructions[i + 2].instr.real() + while i < block.instructions.len() { + let Some(Instruction::BinaryOp { .. }) = block.instructions[i].instr.real() else { + i += 1; + continue; + }; + + let Some(operand_indices) = i + .checked_sub(1) + .and_then(|start| Self::get_const_loading_instr_indices(block, start, 2)) else { i += 1; continue; }; - let op_raw = u32::from(block.instructions[i + 2].arg); + let op_raw = u32::from(block.instructions[i].arg); let Ok(op) = BinOp::try_from(op_raw) else { i += 1; continue; }; - let left = Self::get_const_value_from(&self.metadata, &block.instructions[i]); - let right = Self::get_const_value_from(&self.metadata, &block.instructions[i + 1]); + let left = Self::get_const_value_from( + &self.metadata, + &block.instructions[operand_indices[0]], + ); + let right = Self::get_const_value_from( + &self.metadata, + &block.instructions[operand_indices[1]], + ); let (Some(left_val), Some(right_val)) = (left, right) else { i += 1; @@ -759,20 +943,20 @@ impl CodeInfo { continue; } let (const_idx, _) = self.metadata.consts.insert_full(result_const); - // Replace first instruction with LOAD_CONST result - block.instructions[i].instr = Opcode::LoadConst.into(); + let folded_from_nonliteral_expr = operand_indices + .iter() + .any(|&idx| block.instructions[idx].folded_from_nonliteral_expr); + for &idx in &operand_indices { + set_to_nop(&mut block.instructions[idx]); + block.instructions[idx].location = block.instructions[i].location; + block.instructions[idx].end_location = block.instructions[i].end_location; + } + block.instructions[i].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); block.instructions[i].arg = OpArg::new(const_idx as u32); - // NOP out the second and third instructions - let loc = block.instructions[i].location; - let end_loc = block.instructions[i].end_location; - set_to_nop(&mut block.instructions[i + 1]); - block.instructions[i + 1].location = loc; - block.instructions[i + 1].end_location = end_loc; - set_to_nop(&mut block.instructions[i + 2]); - block.instructions[i + 2].location = loc; - block.instructions[i + 2].end_location = end_loc; - // Don't advance - check if the result can be folded again - // (e.g., 2 ** 31 - 1) + block.instructions[i].folded_from_nonliteral_expr = folded_from_nonliteral_expr; i = i.saturating_sub(1); // re-check with previous instruction } else { i += 1; @@ -781,6 +965,13 @@ impl CodeInfo { } } + fn get_const_value_from_dummy(info: &InstructionInfo) -> Option<()> { + match info.instr.real() { + Some(Instruction::LoadConst { .. } | Instruction::LoadSmallInt { .. }) => Some(()), + _ => None, + } + } + fn get_const_value_from( metadata: &CodeUnitMetadata, info: &InstructionInfo, @@ -806,6 +997,42 @@ impl CodeInfo { op: oparg::BinaryOperator, ) -> Option { use oparg::BinaryOperator as BinOp; + fn eval_complex_binop( + left: Complex, + right: Complex, + op: BinOp, + ) -> Option { + let value = match op { + BinOp::Add => left + right, + BinOp::Subtract => { + let re = left.re - right.re; + let mut im = left.im - right.im; + // Preserve CPython's signed-zero behavior for real-zero + // minus zero-complex expressions such as `0 - 0j`. + if left.re == 0.0 + && left.im == 0.0 + && right.re == 0.0 + && right.im == 0.0 + && !right.im.is_sign_negative() + { + im = -0.0; + } + Complex::new(re, im) + } + BinOp::Multiply => left * right, + BinOp::TrueDivide => { + if right == Complex::new(0.0, 0.0) { + return None; + } + left / right + } + _ => return None, + }; + if !value.re.is_finite() || !value.im.is_finite() { + return None; + } + Some(ConstantData::Complex { value }) + } match (left, right) { (ConstantData::Integer { value: l }, ConstantData::Integer { value: r }) => { let result = match op { @@ -817,6 +1044,18 @@ impl CodeInfo { } l * r } + BinOp::TrueDivide => { + if r.is_zero() { + return None; + } + let l_f = l.to_f64()?; + let r_f = r.to_f64()?; + let result = l_f / r_f; + if !result.is_finite() { + return None; + } + return Some(ConstantData::Float { value: result }); + } BinOp::FloorDivide => { if r.is_zero() { return None; @@ -886,8 +1125,16 @@ impl CodeInfo { return None; } BinOp::Remainder => { - // Float modulo uses fmod() at runtime; Rust arithmetic differs - return None; + if *r == 0.0 { + return None; + } + let mut result = l % r; + if result != 0.0 && (*r < 0.0) != (result < 0.0) { + result += r; + } else if result == 0.0 { + result = 0.0f64.copysign(*r); + } + result } BinOp::Power => l.powf(*r), _ => return None, @@ -914,6 +1161,21 @@ impl CodeInfo { op, ) } + (ConstantData::Integer { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(Complex::new(l.to_f64()?, 0.0), *r, op) + } + (ConstantData::Complex { value: l }, ConstantData::Integer { value: r }) => { + eval_complex_binop(*l, Complex::new(r.to_f64()?, 0.0), op) + } + (ConstantData::Float { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(Complex::new(*l, 0.0), *r, op) + } + (ConstantData::Complex { value: l }, ConstantData::Float { value: r }) => { + eval_complex_binop(*l, Complex::new(*r, 0.0), op) + } + (ConstantData::Complex { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(*l, *r, op) + } // String concatenation and repetition (ConstantData::Str { value: l }, ConstantData::Str { value: r }) if matches!(op, BinOp::Add) => @@ -936,6 +1198,43 @@ impl CodeInfo { value: result.into(), }) } + (ConstantData::Integer { value: n }, ConstantData::Str { value: s }) + if matches!(op, BinOp::Multiply) => + { + let n: usize = n.try_into().ok()?; + if n > 4096 { + return None; + } + let result = s.to_string().repeat(n); + Some(ConstantData::Str { + value: result.into(), + }) + } + (ConstantData::Bytes { value: l }, ConstantData::Bytes { value: r }) + if matches!(op, BinOp::Add) => + { + let mut result = l.clone(); + result.extend_from_slice(r); + Some(ConstantData::Bytes { value: result }) + } + (ConstantData::Bytes { value: b }, ConstantData::Integer { value: n }) + if matches!(op, BinOp::Multiply) => + { + let n: usize = n.try_into().ok()?; + if n > 4096 { + return None; + } + Some(ConstantData::Bytes { value: b.repeat(n) }) + } + (ConstantData::Integer { value: n }, ConstantData::Bytes { value: b }) + if matches!(op, BinOp::Multiply) => + { + let n: usize = n.try_into().ok()?; + if n > 4096 { + return None; + } + Some(ConstantData::Bytes { value: b.repeat(n) }) + } _ => None, } } @@ -944,6 +1243,7 @@ impl CodeInfo { match c { ConstantData::Integer { value } => value.bits() > 4096 * 8, ConstantData::Str { value } => value.len() > 4096, + ConstantData::Bytes { value } => value.len() > 4096, _ => false, } } @@ -962,6 +1262,23 @@ impl CodeInfo { }; let tuple_size = u32::from(instr.arg) as usize; + if block + .instructions + .get(i + 1) + .and_then(|next| next.instr.real()) + .is_some_and(|next| { + matches!( + next, + Instruction::UnpackSequence { .. } + if usize::try_from(u32::from(block.instructions[i + 1].arg)) + .ok() + == Some(tuple_size) + ) + }) + { + i += 1; + continue; + } if tuple_size == 0 { // BUILD_TUPLE 0 → LOAD_CONST () let (const_idx, _) = self.metadata.consts.insert_full(ConstantData::Tuple { @@ -972,18 +1289,22 @@ impl CodeInfo { i += 1; continue; } - if i < tuple_size { + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, tuple_size) + }) else { i += 1; continue; - } + }; - // Check if all preceding instructions are constant-loading - let start_idx = i - tuple_size; let mut elements = Vec::with_capacity(tuple_size); let mut all_const = true; - for j in start_idx..i { + for &j in &operand_indices { let load_instr = &block.instructions[j]; + if load_instr.folded_from_nonliteral_expr { + all_const = false; + break; + } match load_instr.instr.real() { Some(Instruction::LoadConst { .. }) => { let const_idx = u32::from(load_instr.arg) as usize; @@ -1026,7 +1347,7 @@ impl CodeInfo { // Replace preceding LOAD instructions with NOP at the // BUILD_TUPLE location so remove_nops() can eliminate them. let folded_loc = block.instructions[i].location; - for j in start_idx..i { + for &j in &operand_indices { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1053,17 +1374,26 @@ impl CodeInfo { }; let list_size = u32::from(instr.arg) as usize; - if list_size == 0 || i < list_size { + if list_size == 0 { i += 1; continue; } - let start_idx = i - list_size; + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, list_size) + }) else { + i += 1; + continue; + }; let mut elements = Vec::with_capacity(list_size); let mut all_const = true; - for j in start_idx..i { + for &j in &operand_indices { let load_instr = &block.instructions[j]; + if load_instr.folded_from_nonliteral_expr { + all_const = false; + break; + } match load_instr.instr.real() { Some(Instruction::LoadConst { .. }) => { let const_idx = u32::from(load_instr.arg) as usize; @@ -1101,22 +1431,29 @@ impl CodeInfo { let end_loc = block.instructions[i].end_location; let eh = block.instructions[i].except_handler; - // slot[start_idx] → BUILD_LIST 0 - block.instructions[start_idx].instr = Opcode::BuildList.into(); - block.instructions[start_idx].arg = OpArg::new(0); - block.instructions[start_idx].location = folded_loc; - block.instructions[start_idx].end_location = end_loc; - block.instructions[start_idx].except_handler = eh; + let build_idx = operand_indices[0]; + let const_idx_slot = operand_indices[1]; - // slot[start_idx+1] → LOAD_CONST (tuple) - block.instructions[start_idx + 1].instr = Opcode::LoadConst.into(); - block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32); - block.instructions[start_idx + 1].location = folded_loc; - block.instructions[start_idx + 1].end_location = end_loc; - block.instructions[start_idx + 1].except_handler = eh; + block.instructions[build_idx].instr = Instruction::BuildList { + count: Arg::marker(), + } + .into(); + block.instructions[build_idx].arg = OpArg::new(0); + block.instructions[build_idx].location = folded_loc; + block.instructions[build_idx].end_location = end_loc; + block.instructions[build_idx].except_handler = eh; + + block.instructions[const_idx_slot].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); + block.instructions[const_idx_slot].arg = OpArg::new(const_idx as u32); + block.instructions[const_idx_slot].location = folded_loc; + block.instructions[const_idx_slot].end_location = end_loc; + block.instructions[const_idx_slot].except_handler = eh; // NOP the rest - for j in (start_idx + 2)..i { + for &j in &operand_indices[2..] { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1137,6 +1474,22 @@ impl CodeInfo { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { + if matches!( + block.instructions[i].instr.real(), + Some(Instruction::CallIntrinsic1 { func }) + if func.get(block.instructions[i].arg) == IntrinsicFunction1::ListToTuple + ) && matches!( + block + .instructions + .get(i + 1) + .and_then(|instr| instr.instr.real()), + Some(Instruction::GetIter) + ) { + set_to_nop(&mut block.instructions[i]); + i += 2; + continue; + } + let is_build = matches!( block.instructions[i].instr.real(), Some(Instruction::BuildList { .. }) @@ -1186,12 +1539,17 @@ impl CodeInfo { ) { let seq_size = u32::from(block.instructions[i].arg) as usize; - if seq_size != 0 && i >= seq_size { - let start_idx = i - seq_size; + if seq_size != 0 { + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, seq_size) + }) else { + i += 2; + continue; + }; let mut elements = Vec::with_capacity(seq_size); let mut all_const = true; - for j in start_idx..i { + for &j in &operand_indices { match Self::get_const_value_from(&self.metadata, &block.instructions[j]) { Some(constant) => elements.push(constant), @@ -1207,7 +1565,7 @@ impl CodeInfo { let (const_idx, _) = self.metadata.consts.insert_full(const_data); let folded_loc = block.instructions[i].location; - for j in start_idx..i { + for &j in &operand_indices { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1241,17 +1599,26 @@ impl CodeInfo { }; let set_size = u32::from(instr.arg) as usize; - if set_size < 3 || i < set_size { + if set_size < 3 { i += 1; continue; } - let start_idx = i - set_size; + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, set_size) + }) else { + i += 1; + continue; + }; let mut elements = Vec::with_capacity(set_size); let mut all_const = true; - for j in start_idx..i { + for &j in &operand_indices { let load_instr = &block.instructions[j]; + if load_instr.folded_from_nonliteral_expr { + all_const = false; + break; + } match load_instr.instr.real() { Some(Instruction::LoadConst { .. }) => { let const_idx = u32::from(load_instr.arg) as usize; @@ -1282,27 +1649,35 @@ impl CodeInfo { continue; } - // Use FrozenSet constant (stored as Tuple for now) - let const_data = ConstantData::Tuple { elements }; + let const_data = ConstantData::Frozenset { elements }; let (const_idx, _) = self.metadata.consts.insert_full(const_data); let folded_loc = block.instructions[i].location; let end_loc = block.instructions[i].end_location; let eh = block.instructions[i].except_handler; - block.instructions[start_idx].instr = Opcode::BuildSet.into(); - block.instructions[start_idx].arg = OpArg::new(0); - block.instructions[start_idx].location = folded_loc; - block.instructions[start_idx].end_location = end_loc; - block.instructions[start_idx].except_handler = eh; + let build_idx = operand_indices[0]; + let const_idx_slot = operand_indices[1]; - block.instructions[start_idx + 1].instr = Opcode::LoadConst.into(); - block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32); - block.instructions[start_idx + 1].location = folded_loc; - block.instructions[start_idx + 1].end_location = end_loc; - block.instructions[start_idx + 1].except_handler = eh; + block.instructions[build_idx].instr = Instruction::BuildSet { + count: Arg::marker(), + } + .into(); + block.instructions[build_idx].arg = OpArg::new(0); + block.instructions[build_idx].location = folded_loc; + block.instructions[build_idx].end_location = end_loc; + block.instructions[build_idx].except_handler = eh; + + block.instructions[const_idx_slot].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); + block.instructions[const_idx_slot].arg = OpArg::new(const_idx as u32); + block.instructions[const_idx_slot].location = folded_loc; + block.instructions[const_idx_slot].end_location = end_loc; + block.instructions[const_idx_slot].except_handler = eh; - for j in (start_idx + 2)..i { + for &j in &operand_indices[2..] { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1366,6 +1741,8 @@ impl CodeInfo { /// intervening swappable stores to one of the same variables. Do not /// cross line-number boundaries (user-visible name bindings). fn apply_static_swaps(&mut self) { + const VISITED: i32 = -1; + /// Instruction classes that are safe to reorder around SWAP. fn is_swappable(instr: &AnyInstruction) -> bool { matches!( @@ -1410,59 +1787,111 @@ impl CodeInfo { } } - for block in &mut self.blocks { - let instructions = &mut block.instructions; - let len = instructions.len(); - // Walk forward; for each SWAP attempt elimination. - let mut i = 0; - while i < len { - let swap_arg = match instructions[i].instr { - AnyInstruction::Real(Instruction::Swap { .. }) => { - u32::from(instructions[i].arg) + fn optimize_swap_block(instructions: &mut [InstructionInfo]) { + let mut i = 0usize; + while i < instructions.len() { + let AnyInstruction::Real(Instruction::Swap { .. }) = instructions[i].instr else { + i += 1; + continue; + }; + + let mut len = 0usize; + let mut depth = 0usize; + let mut more = false; + while i + len < instructions.len() { + let info = &instructions[i + len]; + match info.instr.real() { + Some(Instruction::Swap { .. }) => { + let oparg = u32::from(info.arg) as usize; + depth = depth.max(oparg); + more |= len > 0; + len += 1; + } + Some(Instruction::Nop) => { + len += 1; + } + _ => break, } - _ => { - i += 1; + } + + if !more { + i += len.max(1); + continue; + } + + let mut stack: Vec = (0..depth as i32).collect(); + for info in &instructions[i..i + len] { + if matches!(info.instr.real(), Some(Instruction::Swap { .. })) { + let oparg = u32::from(info.arg) as usize; + stack.swap(0, oparg - 1); + } + } + + let mut current = len as isize - 1; + for slot in 0..depth { + if stack[slot] == VISITED || stack[slot] == slot as i32 { + continue; + } + let mut j = slot; + loop { + if j != 0 { + let out = &mut instructions[i + current as usize]; + out.instr = Opcode::Swap.into(); + out.arg = OpArg::new((j + 1) as u32); + out.target = BlockIdx::NULL; + current -= 1; + } + if stack[j] == VISITED { + debug_assert_eq!(j, slot); + break; + } + let next_j = stack[j] as usize; + stack[j] = VISITED; + j = next_j; + } + } + while current >= 0 { + set_to_nop(&mut instructions[i + current as usize]); + current -= 1; + } + i += len; + } + } + + fn apply_from(instructions: &mut [InstructionInfo], mut i: isize) { + while i >= 0 { + let idx = i as usize; + let swap_arg = match instructions[idx].instr.real() { + Some(Instruction::Swap { .. }) => u32::from(instructions[idx].arg), + Some(Instruction::Nop) + | Some(Instruction::PopTop | Instruction::StoreFast { .. }) => { + i -= 1; continue; } + _ => return, }; - // SWAP oparg < 2 is a no-op; the compiler should not emit - // these, but be defensive. + if swap_arg < 2 { - i += 1; - continue; + return; } - // Find first swappable after SWAP (lineno = -1 initially). - let Some(j) = next_swappable(instructions, i, -1) else { - i += 1; - continue; + + let Some(j) = next_swappable(instructions, idx, -1) else { + return; }; let lineno = instructions[j].location.line.get() as i32; - // Walk (swap_arg - 1) more swappable instructions, with - // lineno constraint. let mut k = j; - let mut ok = true; for _ in 1..swap_arg { - match next_swappable(instructions, k, lineno) { - Some(next) => k = next, - None => { - ok = false; - break; - } - } - } - if !ok { - i += 1; - continue; + let Some(next) = next_swappable(instructions, k, lineno) else { + return; + }; + k = next; } - // Conflict check: if either j or k is a STORE_FAST, no - // intervening store may target the same variable, and - // they must not target the same variable themselves. + let store_j = stores_to(&instructions[j]); let store_k = stores_to(&instructions[k]); if store_j.is_some() || store_k.is_some() { if store_j == store_k { - i += 1; - continue; + return; } let conflict = instructions[(j + 1)..k].iter().any(|info| { if let Some(store_idx) = stores_to(info) { @@ -1472,15 +1901,27 @@ impl CodeInfo { } }); if conflict { - i += 1; - continue; + return; } } - // Safe to reorder. SWAP -> NOP, swap j and k. - instructions[i].instr = Opcode::Nop.into(); - instructions[i].arg = OpArg::new(0); + + instructions[idx].instr = Opcode::Nop.into(); + instructions[idx].arg = OpArg::new(0); instructions.swap(j, k); - i += 1; + i -= 1; + } + } + + for block in &mut self.blocks { + optimize_swap_block(&mut block.instructions); + let len = block.instructions.len(); + for i in 0..len { + if matches!( + block.instructions[i].instr.real(), + Some(Instruction::Swap { .. }) + ) { + apply_from(&mut block.instructions, i as isize); + } } } } @@ -1488,14 +1929,12 @@ impl CodeInfo { /// Eliminate dead stores in STORE_FAST sequences (apply_static_swaps). /// /// In sequences of consecutive STORE_FAST instructions (from tuple unpacking), - /// if the same variable is stored to more than once, only the first store - /// (which gets TOS — the rightmost value) matters. Later stores to the - /// same variable are dead and replaced with POP_TOP. - /// Simplified apply_static_swaps (CPython flowgraph.c): - /// In STORE_FAST sequences that follow UNPACK_SEQUENCE / UNPACK_EX, - /// replace duplicate stores to the same variable with POP_TOP. - /// UNPACK pushes values so stores execute left-to-right; the LAST - /// store to a variable carries the final value, earlier ones are dead. + /// only collapse directly adjacent duplicate targets. + /// + /// CPython preserves non-adjacent duplicates such as `_, expr, _` so the + /// store layout still reflects the original unpack order. Replacing the + /// first `_` with POP_TOP there changes the emitted superinstructions and + /// bytecode shape even though the final value is the same. fn eliminate_dead_stores(&mut self) { for block in &mut self.blocks { let instructions = &mut block.instructions; @@ -1523,18 +1962,18 @@ impl CodeInfo { run_end += 1; } if run_end - run_start >= 2 { - // Pass 1: find the LAST occurrence of each variable - let mut last_occurrence = std::collections::HashMap::new(); - for (j, instr) in instructions[run_start..run_end].iter().enumerate() { - last_occurrence.insert(u32::from(instr.arg), j); - } - // Pass 2: non-last stores to the same variable are dead - for (j, instr) in instructions[run_start..run_end].iter_mut().enumerate() { - let idx = u32::from(instr.arg); - if last_occurrence[&idx] != j { + let mut j = run_start; + while j < run_end { + let arg = u32::from(instructions[j].arg); + let mut group_end = j + 1; + while group_end < run_end && u32::from(instructions[group_end].arg) == arg { + group_end += 1; + } + for instr in &mut instructions[j..group_end.saturating_sub(1)] { instr.instr = Opcode::PopTop.into(); instr.arg = OpArg::new(0); } + j = group_end; } } i = run_end.max(i + 1); @@ -1544,11 +1983,35 @@ impl CodeInfo { /// Peephole optimization: combine consecutive instructions into super-instructions fn peephole_optimize(&mut self) { + let const_truthiness = + |instr: Instruction, arg: OpArg, metadata: &CodeUnitMetadata| match instr { + Instruction::LoadConst { consti } => { + let constant = &metadata.consts[consti.get(arg).as_usize()]; + Some(match constant { + ConstantData::Tuple { elements } => !elements.is_empty(), + ConstantData::Integer { value } => !value.is_zero(), + ConstantData::Float { value } => *value != 0.0, + ConstantData::Complex { value } => value.re != 0.0 || value.im != 0.0, + ConstantData::Boolean { value } => *value, + ConstantData::Str { value } => !value.is_empty(), + ConstantData::Bytes { value } => !value.is_empty(), + ConstantData::Code { .. } => true, + ConstantData::Slice { .. } => true, + ConstantData::Frozenset { elements } => !elements.is_empty(), + ConstantData::None => false, + ConstantData::Ellipsis => true, + }) + } + Instruction::LoadSmallInt { i } => Some(i.get(arg) != 0), + _ => None, + }; for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; let next = &block.instructions[i + 1]; + let curr_arg = curr.arg; + let next_arg = next.arg; // Only combine if both are real instructions (not pseudo) let (Some(curr_instr), Some(next_instr)) = (curr.instr.real(), next.instr.real()) @@ -1557,65 +2020,136 @@ impl CodeInfo { continue; }; + if let Some(is_true) = const_truthiness(curr_instr, curr.arg, &self.metadata) { + let jump_if_true = match next_instr { + Instruction::PopJumpIfTrue { .. } => Some(true), + Instruction::PopJumpIfFalse { .. } => Some(false), + _ => None, + }; + if let Some(jump_if_true) = jump_if_true { + let target = match next_instr { + Instruction::PopJumpIfTrue { delta } + | Instruction::PopJumpIfFalse { delta } => delta.get(next.arg), + _ => unreachable!(), + }; + set_to_nop(&mut block.instructions[i]); + if is_true == jump_if_true { + block.instructions[i + 1].instr = PseudoInstruction::Jump { + delta: Arg::marker(), + } + .into(); + block.instructions[i + 1].arg = OpArg::new(u32::from(target)); + } else { + set_to_nop(&mut block.instructions[i + 1]); + } + i += 1; + continue; + } + } + + if let Instruction::LoadConst { consti } = curr_instr { + let constant = &self.metadata.consts[consti.get(curr_arg).as_usize()]; + if matches!(constant, ConstantData::None) + && let Instruction::IsOp { invert } = next_instr + { + let mut jump_idx = i + 2; + if jump_idx >= block.instructions.len() { + i += 1; + continue; + } + + if matches!( + block.instructions[jump_idx].instr.real(), + Some(Instruction::ToBool) + ) { + set_to_nop(&mut block.instructions[jump_idx]); + jump_idx += 1; + if jump_idx >= block.instructions.len() { + i += 1; + continue; + } + } + + let Some(jump_instr) = block.instructions[jump_idx].instr.real() else { + i += 1; + continue; + }; + + let mut invert = matches!( + invert.get(next_arg), + rustpython_compiler_core::bytecode::Invert::Yes + ); + let delta = match jump_instr { + Instruction::PopJumpIfFalse { delta } => { + invert = !invert; + delta.get(block.instructions[jump_idx].arg) + } + Instruction::PopJumpIfTrue { delta } => { + delta.get(block.instructions[jump_idx].arg) + } + _ => { + i += 1; + continue; + } + }; + + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + block.instructions[jump_idx].instr = if invert { + Instruction::PopJumpIfNotNone { + delta: Arg::marker(), + } + } else { + Instruction::PopJumpIfNone { + delta: Arg::marker(), + } + } + .into(); + block.instructions[jump_idx].arg = OpArg::new(u32::from(delta)); + i = jump_idx; + continue; + } + } + if matches!( - next_instr.into(), - Opcode::PopJumpIfFalse | Opcode::PopJumpIfTrue - ) && matches!(curr_instr.into(), Opcode::CompareOp) + curr_instr, + Instruction::LoadConst { .. } | Instruction::LoadSmallInt { .. } + ) && matches!(next_instr, Instruction::PopTop) { - block.instructions[i].arg = OpArg::new( - u32::from(block.instructions[i].arg) | oparg::COMPARE_OP_BOOL_MASK, - ); + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + i += 1; + continue; + } + + if matches!(curr_instr, Instruction::Copy { i } if i.get(curr.arg) == 1) + && matches!(next_instr, Instruction::PopTop) + { + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); i += 1; continue; } let combined = { match (curr_instr, next_instr) { - // LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16) - (Instruction::LoadFast { .. }, Instruction::LoadFast { .. }) => { - let line1 = curr.location.line.get() as i32; - let line2 = next.location.line.get() as i32; - if line1 > 0 && line2 > 0 && line1 != line2 { - None - } else { - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - Some((Opcode::LoadFastLoadFast.into(), OpArg::new(packed))) - } else { - None - } - } - } - // StoreFast + StoreFast -> StoreFastStoreFast (if both indices < 16) - // Dead store elimination: if both store to the same variable, - // the first store is dead. Replace it with POP_TOP (like - // apply_static_swaps in CPython's flowgraph.c). - (Instruction::StoreFast { .. }, Instruction::StoreFast { .. }) => { - let line1 = curr.location.line.get() as i32; - let line2 = next.location.line.get() as i32; - if line1 > 0 && line2 > 0 && line1 != line2 { - None - } else { - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - Some((Opcode::StoreFastStoreFast.into(), OpArg::new(packed))) - } else { - None - } - } - } // Note: StoreFast + LoadFast → StoreFastLoadFast is done in a - // separate pass AFTER optimize_load_fast_borrow, because CPython - // only combines STORE_FAST + LOAD_FAST (not LOAD_FAST_BORROW). - (Instruction::LoadConst { consti }, Instruction::ToBool) => { - let consti = consti.get(curr.arg); - let constant = &self.metadata.consts[consti.as_usize()]; - if let ConstantData::Boolean { .. } = constant { - Some((curr_instr, OpArg::from(consti.as_u32()))) + // later pass aligned with CPython insert_superinstructions(). + (Instruction::LoadConst { .. }, Instruction::ToBool) + | (Instruction::LoadSmallInt { .. }, Instruction::ToBool) => { + if let Some(value) = + const_truthiness(curr_instr, curr.arg, &self.metadata) + { + let (const_idx, _) = self + .metadata + .consts + .insert_full(ConstantData::Boolean { value }); + Some(( + Instruction::LoadConst { + consti: Arg::marker(), + }, + OpArg::new(const_idx as u32), + )) } else { None } @@ -1686,6 +2220,39 @@ impl CodeInfo { } } + fn remove_redundant_const_pop_top_pairs(&mut self) { + for block in &mut self.blocks { + let mut i = 0; + while i + 1 < block.instructions.len() { + let curr = &block.instructions[i]; + let next = &block.instructions[i + 1]; + let Some(curr_instr) = curr.instr.real() else { + i += 1; + continue; + }; + let Some(next_instr) = next.instr.real() else { + i += 1; + continue; + }; + + let redundant = matches!( + (curr_instr, next_instr), + (Instruction::LoadConst { .. }, Instruction::PopTop) + | (Instruction::LoadSmallInt { .. }, Instruction::PopTop) + ) || matches!(curr_instr, Instruction::Copy { i } if i.get(curr.arg) == 1) + && matches!(next_instr, Instruction::PopTop); + + if redundant { + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + i += 2; + } else { + i += 1; + } + } + } + } + /// Convert LOAD_CONST for small integers to LOAD_SMALL_INT /// maybe_instr_make_load_smallint fn convert_to_load_small_int(&mut self) { @@ -1787,52 +2354,111 @@ impl CodeInfo { let mut prev_line = None; block.instructions.retain(|ins| { if matches!(ins.instr.real(), Some(Instruction::Nop)) { - let line = ins.location.line; + let line = ins.location.line.get() as i32; if prev_line == Some(line) { return false; } } - prev_line = Some(ins.location.line); + prev_line = Some(instruction_lineno(ins)); true }); } } - /// Optimize LOAD_FAST to LOAD_FAST_BORROW where safe. - /// - /// insert_superinstructions (flowgraph.c): Combine STORE_FAST + LOAD_FAST → - /// STORE_FAST_LOAD_FAST. Currently disabled pending VM stack null investigation. - #[allow(dead_code)] - fn combine_store_fast_load_fast(&mut self) { + /// insert_superinstructions (flowgraph.c): combine a narrow subset of + /// STORE_FAST + LOAD_FAST patterns that CPython uses in comprehension loop + /// headers. Keeping this scoped avoids reintroducing earlier mismatches in + /// non-loop code while we continue aligning the surrounding borrow rules. + fn insert_superinstructions(&mut self) { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; - let next = &block.instructions[i + 1]; - let (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) = - (curr.instr.real(), next.instr.real()) - else { + let line = curr.location.line; + + let mut j = i + 1; + while j < block.instructions.len() + && matches!(block.instructions[j].instr.real(), Some(Instruction::Nop)) + && block.instructions[j].location.line == line + { + j += 1; + } + if j >= block.instructions.len() { i += 1; continue; - }; - // Skip if instructions are on different lines (matching make_super_instruction) - let line1 = curr.location.line; - let line2 = next.location.line; - if line1 != line2 { + } + + let next = &block.instructions[j]; + if next.location.line != line { i += 1; continue; } - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - block.instructions[i].instr = Opcode::StoreFastLoadFast.into(); - block.instructions[i].arg = OpArg::new(packed); - // Replace second instruction with NOP (CPython: INSTR_SET_OP0(inst2, NOP)) - set_to_nop(&mut block.instructions[i + 1]); - i += 2; // skip the NOP - } else { - i += 1; + + let prev_real = block.instructions[..i] + .iter() + .rev() + .find_map(|info| info.instr.real()); + let next_real = block.instructions[(j + 1)..] + .iter() + .find_map(|info| info.instr.real()); + + match (curr.instr.real(), next.instr.real()) { + (Some(Instruction::LoadFast { .. }), Some(Instruction::LoadFast { .. })) => { + let idx1 = u32::from(curr.arg); + let idx2 = u32::from(next.arg); + if idx1 >= 16 || idx2 >= 16 { + i += 1; + continue; + } + let packed = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + block.instructions.drain(i + 1..=j); + } + ( + Some(Instruction::StoreFast { .. }), + Some(Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. }), + ) => { + if self.flags.contains(CodeFlags::GENERATOR) + && matches!(prev_real, Some(Instruction::ForIter { .. })) + && !matches!(next_real, Some(Instruction::ToBool)) + { + i += 1; + continue; + } + let store_idx = u32::from(curr.arg); + let load_idx = u32::from(next.arg); + if store_idx >= 16 || load_idx >= 16 { + i += 1; + continue; + } + let packed = (store_idx << 4) | load_idx; + block.instructions[i].instr = Instruction::StoreFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + block.instructions.drain(i + 1..=j); + } + (Some(Instruction::StoreFast { .. }), Some(Instruction::StoreFast { .. })) => { + let idx1 = u32::from(curr.arg); + let idx2 = u32::from(next.arg); + if idx1 >= 16 || idx2 >= 16 { + i += 1; + continue; + } + let packed = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::StoreFastStoreFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + block.instructions.drain(i + 1..=j); + } + _ => i += 1, } } } @@ -1841,82 +2467,464 @@ impl CodeInfo { fn optimize_load_fast_borrow(&mut self) { // NOT_LOCAL marker: instruction didn't come from a LOAD_FAST const NOT_LOCAL: usize = usize::MAX; + const DUMMY_INSTR: isize = -1; + const SUPPORT_KILLED: u8 = 1; + const STORED_AS_LOCAL: u8 = 2; + const REF_UNCONSUMED: u8 = 4; + + #[derive(Clone, Copy)] + struct AbstractRef { + instr: isize, + local: usize, + } - for block in &mut self.blocks { - if block.instructions.is_empty() { - continue; + fn push_ref(refs: &mut Vec, instr: isize, local: usize) { + refs.push(AbstractRef { instr, local }); + } + + fn pop_ref(refs: &mut Vec) -> Option { + refs.pop() + } + + fn at_ref(refs: &[AbstractRef], idx: usize) -> Option { + refs.get(idx).copied() + } + + fn swap_top(refs: &mut [AbstractRef], depth: usize) { + let top = refs.len() - 1; + let other = refs.len() - depth; + refs.swap(top, other); + } + + fn kill_local(instr_flags: &mut [u8], refs: &[AbstractRef], local: usize) { + for r in refs.iter().copied().filter(|r| r.local == local) { + debug_assert!(r.instr >= 0); + instr_flags[r.instr as usize] |= SUPPORT_KILLED; } + } - // Track which instructions' outputs are still on stack at block end - // For each instruction, we track if its pushed value(s) are unconsumed - let mut unconsumed = vec![false; block.instructions.len()]; + fn store_local(instr_flags: &mut [u8], refs: &[AbstractRef], local: usize, r: AbstractRef) { + kill_local(instr_flags, refs, local); + if r.instr != DUMMY_INSTR { + instr_flags[r.instr as usize] |= STORED_AS_LOCAL; + } + } - // Simulate stack: each entry is the instruction index that pushed it - // (or NOT_LOCAL if not from LOAD_FAST/LOAD_FAST_LOAD_FAST). - // - // CPython (flowgraph.c optimize_load_fast) pre-fills the stack with - // dummy refs for values inherited from predecessor blocks. We take - // the simpler approach of aborting the optimisation for the whole - // block on stack underflow. - let mut stack: Vec = Vec::new(); - let mut underflow = false; + fn decode_packed_fast_locals(arg: OpArg) -> (usize, usize) { + let packed = u32::from(arg); + (((packed >> 4) & 0xF) as usize, (packed & 0xF) as usize) + } - for (i, info) in block.instructions.iter().enumerate() { - let Some(instr) = info.instr.real() else { - continue; - }; + fn push_block( + worklist: &mut Vec, + visited: &mut [bool], + blocks: &[Block], + source: BlockIdx, + target: BlockIdx, + start_depth: usize, + ) { + let expected = blocks[target.idx()].start_depth.map(|depth| depth as usize); + if expected != Some(start_depth) { + debug_assert!( + expected == Some(start_depth), + "optimize_load_fast_borrow start_depth mismatch: source={source:?} target={target:?} expected={expected:?} actual={:?} source_last={:?} target_instrs={:?}", + Some(start_depth), + blocks[source.idx()] + .instructions + .last() + .and_then(|info| info.instr.real()), + blocks[target.idx()] + .instructions + .iter() + .map(|info| info.instr) + .collect::>(), + ); + return; + } + if !visited[target.idx()] { + visited[target.idx()] = true; + worklist.push(target); + } + } - let stack_effect_info = instr.stack_effect_info(info.arg.into()); - let (pushes, pops) = (stack_effect_info.pushed(), stack_effect_info.popped()); + let mut visited = vec![false; self.blocks.len()]; + let mut worklist = vec![BlockIdx(0)]; + visited[0] = true; - // Pop values from stack - for _ in 0..pops { - if stack.pop().is_none() { - // Stack underflow — block receives values from a predecessor. - // Abort optimisation for the entire block. - underflow = true; - break; + while let Some(block_idx) = worklist.pop() { + let block = &self.blocks[block_idx]; + + let mut instr_flags = vec![0u8; block.instructions.len()]; + let start_depth = block.start_depth.unwrap_or(0) as usize; + let mut refs = Vec::with_capacity(block.instructions.len() + start_depth + 2); + for _ in 0..start_depth { + push_ref(&mut refs, DUMMY_INSTR, NOT_LOCAL); + } + + for (i, info) in block.instructions.iter().enumerate() { + let instr = info.instr; + let arg_u32 = u32::from(info.arg); + + match instr { + AnyInstruction::Real(Instruction::DeleteFast { var_num }) => { + kill_local(&mut instr_flags, &refs, usize::from(var_num.get(info.arg))); + } + AnyInstruction::Real(Instruction::LoadFast { var_num }) => { + push_ref(&mut refs, i as isize, usize::from(var_num.get(info.arg))); + } + AnyInstruction::Real(Instruction::LoadFastAndClear { var_num }) => { + let local = usize::from(var_num.get(info.arg)); + kill_local(&mut instr_flags, &refs, local); + push_ref(&mut refs, i as isize, local); + } + AnyInstruction::Real(Instruction::LoadFastLoadFast { .. }) => { + let (local1, local2) = decode_packed_fast_locals(info.arg); + push_ref(&mut refs, i as isize, local1); + push_ref(&mut refs, i as isize, local2); + } + AnyInstruction::Real(Instruction::StoreFast { var_num }) => { + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local( + &mut instr_flags, + &refs, + usize::from(var_num.get(info.arg)), + r, + ); + } + AnyInstruction::Pseudo(PseudoInstruction::StoreFastMaybeNull { var_num }) => { + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, var_num.get(info.arg) as usize, r); + } + AnyInstruction::Real(Instruction::StoreFastLoadFast { .. }) => { + let (store_local_idx, load_local_idx) = decode_packed_fast_locals(info.arg); + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, store_local_idx, r); + push_ref(&mut refs, i as isize, load_local_idx); + } + AnyInstruction::Real(Instruction::StoreFastStoreFast { .. }) => { + let (local1, local2) = decode_packed_fast_locals(info.arg); + let Some(r1) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, local1, r1); + let Some(r2) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, local2, r2); + } + AnyInstruction::Real(Instruction::Copy { i: _ }) => { + let depth = arg_u32 as usize; + if depth == 0 || refs.len() < depth { + continue; + } + let r = at_ref(&refs, refs.len() - depth).expect("copy index in bounds"); + push_ref(&mut refs, r.instr, r.local); + } + AnyInstruction::Real(Instruction::Swap { i: _ }) => { + let depth = arg_u32 as usize; + if depth < 2 || refs.len() < depth { + continue; + } + swap_top(&mut refs, depth); + } + AnyInstruction::Real( + Instruction::FormatSimple + | Instruction::GetANext + | Instruction::GetLen + | Instruction::GetYieldFromIter + | Instruction::ImportFrom { .. } + | Instruction::MatchKeys + | Instruction::MatchMapping + | Instruction::MatchSequence + | Instruction::WithExceptStart, + ) => { + let effect = instr.stack_effect_info(arg_u32); + let net_pushed = effect.pushed() as isize - effect.popped() as isize; + debug_assert!(net_pushed >= 0); + for _ in 0..net_pushed { + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + } + AnyInstruction::Real( + Instruction::DictMerge { .. } + | Instruction::DictUpdate { .. } + | Instruction::ListAppend { .. } + | Instruction::ListExtend { .. } + | Instruction::MapAdd { .. } + | Instruction::Reraise { .. } + | Instruction::SetAdd { .. } + | Instruction::SetUpdate { .. }, + ) => { + let effect = instr.stack_effect_info(arg_u32); + let net_popped = effect.popped() as isize - effect.pushed() as isize; + debug_assert!(net_popped > 0); + for _ in 0..net_popped { + let _ = pop_ref(&mut refs); + } + } + AnyInstruction::Real( + Instruction::EndSend | Instruction::SetFunctionAttribute { .. }, + ) => { + let Some(tos) = pop_ref(&mut refs) else { + continue; + }; + let _ = pop_ref(&mut refs); + push_ref(&mut refs, tos.instr, tos.local); + } + AnyInstruction::Real(Instruction::CheckExcMatch) => { + let _ = pop_ref(&mut refs); + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + AnyInstruction::Real(Instruction::ForIter { .. }) => { + let target = info.target; + if target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + target, + refs.len() + 1, + ); + } + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + AnyInstruction::Real(Instruction::LoadAttr { .. }) => { + let Some(self_ref) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + if arg_u32 & 1 != 0 { + push_ref(&mut refs, self_ref.instr, self_ref.local); + } + } + AnyInstruction::Real(Instruction::LoadSuperAttr { .. }) => { + let _ = pop_ref(&mut refs); + let _ = pop_ref(&mut refs); + let Some(self_ref) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + if arg_u32 & 1 != 0 { + push_ref(&mut refs, self_ref.instr, self_ref.local); + } + } + AnyInstruction::Real( + Instruction::LoadSpecial { .. } | Instruction::PushExcInfo, + ) => { + let Some(tos) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + push_ref(&mut refs, tos.instr, tos.local); + } + AnyInstruction::Real(Instruction::Send { .. }) => { + let target = info.target; + if target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + target, + refs.len(), + ); + } + let _ = pop_ref(&mut refs); + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + _ => { + let effect = instr.stack_effect_info(arg_u32); + let num_popped = effect.popped() as usize; + let num_pushed = effect.pushed() as usize; + let target = info.target; + if target != BlockIdx::NULL { + let target_depth = refs + .len() + .saturating_sub(num_popped) + .saturating_add(num_pushed); + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + target, + target_depth, + ); + } + if !instr.is_block_push() { + for _ in 0..num_popped { + let _ = pop_ref(&mut refs); + } + for _ in 0..num_pushed { + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + } } } - if underflow { - break; - } + } - // Push values to stack with source instruction index - let source = match instr.into() { - Opcode::LoadFast | Opcode::LoadFastLoadFast => i, - _ => NOT_LOCAL, - }; - for _ in 0..pushes { - stack.push(source); + let next = block.next; + if next != BlockIdx::NULL + && block.instructions.last().is_none_or(|term| { + !term.instr.is_unconditional_jump() && !term.instr.is_scope_exit() + }) + { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + next, + refs.len(), + ); + } + + for r in refs { + if r.instr != DUMMY_INSTR { + instr_flags[r.instr as usize] |= REF_UNCONSUMED; } } - if underflow { + let block = &mut self.blocks[block_idx]; + if block.disable_load_fast_borrow { continue; } + for (i, info) in block.instructions.iter_mut().enumerate() { + if instr_flags[i] != 0 { + continue; + } + match info.instr.real() { + Some(Instruction::LoadFast { .. }) => { + info.instr = Instruction::LoadFastBorrow { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastLoadFast { .. }) => { + info.instr = Instruction::LoadFastBorrowLoadFastBorrow { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + } + + fn compute_load_fast_start_depths(&mut self) { + fn stackdepth_push( + stack: &mut Vec, + start_depths: &mut [u32], + target: BlockIdx, + depth: u32, + ) { + let idx = target.idx(); + let block_depth = &mut start_depths[idx]; + debug_assert!( + *block_depth == u32::MAX || *block_depth == depth, + "Invalid CFG, inconsistent optimize_load_fast stackdepth for block {:?}: existing={}, new={}", + target, + *block_depth, + depth, + ); + if *block_depth == u32::MAX { + *block_depth = depth; + stack.push(target); + } + } + + let mut stack = Vec::with_capacity(self.blocks.len()); + let mut start_depths = vec![u32::MAX; self.blocks.len()]; + stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0); - // Mark instructions whose values remain on stack at block end - for &src in &stack { - if src != NOT_LOCAL { - unconsumed[src] = true; + 'process_blocks: while let Some(block_idx) = stack.pop() { + let mut depth = start_depths[block_idx.idx()]; + let block = &self.blocks[block_idx]; + for ins in &block.instructions { + let instr = &ins.instr; + let effect = instr.stack_effect(ins.arg.into()); + let new_depth = depth.saturating_add_signed(effect); + if ins.target != BlockIdx::NULL { + let jump_effect = instr.stack_effect_jump(ins.arg.into()); + let target_depth = depth.saturating_add_signed(jump_effect); + stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth); } + depth = new_depth; + if instr.is_scope_exit() || instr.is_unconditional_jump() { + continue 'process_blocks; + } + } + if block.next != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, block.next, depth); } + } - // Convert LOAD_FAST to LOAD_FAST_BORROW where value is fully consumed - for (i, info) in block.instructions.iter_mut().enumerate() { - if unconsumed[i] { + for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) { + block.start_depth = (start_depth != u32::MAX).then_some(start_depth); + } + } + + fn deoptimize_borrow_for_handler_return_paths(&mut self) { + for block in &mut self.blocks { + let len = block.instructions.len(); + for i in 0..len { + let Some(Instruction::LoadFastBorrow { .. }) = block.instructions[i].instr.real() + else { + continue; + }; + let tail = &block.instructions[i + 1..]; + if tail.len() < 3 { continue; } - let Some(instr) = info.instr.real() else { + if !matches!(tail[0].instr.real(), Some(Instruction::Swap { .. })) { continue; - }; - match instr.into() { - Opcode::LoadFast => { - info.instr = Opcode::LoadFastBorrow.into(); + } + if !matches!(tail[1].instr.real(), Some(Instruction::PopExcept)) { + continue; + } + if !matches!(tail[2].instr.real(), Some(Instruction::ReturnValue)) { + continue; + } + block.instructions[i].instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + } + } + + fn deoptimize_borrow_after_push_exc_info(&mut self) { + for block in &mut self.blocks { + let mut in_exception_state = false; + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::PushExcInfo) => { + in_exception_state = true; + } + Some(Instruction::PopExcept) | Some(Instruction::Reraise { .. }) => { + in_exception_state = false; + } + Some(Instruction::LoadFastBorrow { .. }) if in_exception_state => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); } - Opcode::LoadFastLoadFast => { - info.instr = Opcode::LoadFastBorrowLoadFastBorrow.into(); + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) + if in_exception_state => + { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); } _ => {} } @@ -1924,12 +2932,205 @@ impl CodeInfo { } } + fn deoptimize_borrow_for_match_keys_attr(&mut self) { + let Some(key_name_idx) = self.metadata.names.get_index_of("KEY") else { + return; + }; + + let mut to_deopt = Vec::new(); + for block_idx in 0..self.blocks.len() { + let block = &self.blocks[block_idx]; + let len = block.instructions.len(); + for i in 0..len { + let Some(Instruction::LoadFastBorrow { .. }) = block.instructions[i].instr.real() + else { + continue; + }; + let Some(Instruction::LoadAttr { namei }) = block + .instructions + .get(i + 1) + .and_then(|info| info.instr.real()) + else { + continue; + }; + let load_attr = namei.get(block.instructions[i + 1].arg); + if load_attr.is_method() || load_attr.name_idx() as usize != key_name_idx { + continue; + } + + let mut saw_build_tuple = false; + let mut saw_match_keys = false; + let mut scan_block_idx = block_idx; + let mut scan_start = i + 2; + loop { + let scan_block = &self.blocks[scan_block_idx]; + for info in scan_block.instructions.iter().skip(scan_start) { + match info.instr.real() { + Some( + Instruction::LoadConst { .. } + | Instruction::LoadSmallInt { .. } + | Instruction::LoadFast { .. } + | Instruction::LoadFastBorrow { .. } + | Instruction::LoadAttr { .. } + | Instruction::Nop, + ) => {} + Some(Instruction::BuildTuple { .. }) => saw_build_tuple = true, + Some(Instruction::MatchKeys) => { + saw_match_keys = true; + break; + } + _ => { + saw_build_tuple = false; + break; + } + } + } + if saw_match_keys { + break; + } + let Some(last) = scan_block.instructions.last() else { + break; + }; + if scan_block.next == BlockIdx::NULL + || last.instr.is_scope_exit() + || last.instr.is_unconditional_jump() + || last.target != BlockIdx::NULL + { + break; + } + scan_block_idx = scan_block.next.idx(); + scan_start = 0; + } + + if saw_build_tuple && saw_match_keys { + to_deopt.push((block_idx, i)); + } + } + } + + for (block_idx, instr_idx) in to_deopt { + self.blocks[block_idx].instructions[instr_idx].instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + } + + fn deoptimize_store_fast_store_fast_after_cleanup(&mut self) { + fn last_real_instr(block: &Block) -> Option { + block + .instructions + .iter() + .rev() + .find_map(|info| info.instr.real()) + } + + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + for (pred_idx, block) in self.blocks.iter().enumerate() { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(BlockIdx(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx(pred_idx as u32)); + } + } + } + + let starts_after_cleanup: Vec = predecessors + .iter() + .map(|predecessor_blocks| { + !predecessor_blocks.is_empty() + && predecessor_blocks.iter().copied().all(|pred_idx| { + matches!( + last_real_instr(&self.blocks[pred_idx]), + Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) + ) + }) + }) + .collect(); + + for (block_idx, block) in self.blocks.iter_mut().enumerate() { + let mut new_instructions = Vec::with_capacity(block.instructions.len()); + let mut in_restore_prefix = starts_after_cleanup[block_idx]; + for (i, info) in block.instructions.iter().copied().enumerate() { + if !in_restore_prefix + && matches!( + info.instr.real(), + Some( + Instruction::StoreFast { .. } | Instruction::StoreFastStoreFast { .. } + ) + ) + && !new_instructions.is_empty() + && new_instructions.iter().all(|prev: &InstructionInfo| { + matches!( + prev.instr.real(), + Some(Instruction::Swap { .. }) | Some(Instruction::PopTop) + ) + }) + { + in_restore_prefix = true; + } + let expand = matches!( + info.instr.real(), + Some(Instruction::StoreFastStoreFast { .. }) + ) && (new_instructions.last().is_some_and( + |prev: &InstructionInfo| { + matches!( + prev.instr.real(), + Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) + ) + }, + ) || (i == 0 && starts_after_cleanup[block_idx]) + || in_restore_prefix); + + if expand { + let Some(Instruction::StoreFastStoreFast { var_nums }) = info.instr.real() + else { + unreachable!(); + }; + let packed = var_nums.get(info.arg); + let (idx1, idx2) = packed.indexes(); + + let mut first = info; + first.instr = Instruction::StoreFast { + var_num: Arg::marker(), + } + .into(); + first.arg = OpArg::new(u32::from(idx1)); + new_instructions.push(first); + + let mut second = info; + second.instr = Instruction::StoreFast { + var_num: Arg::marker(), + } + .into(); + second.arg = OpArg::new(u32::from(idx2)); + new_instructions.push(second); + continue; + } + + in_restore_prefix &= + matches!(info.instr.real(), Some(Instruction::StoreFast { .. })); + new_instructions.push(info); + } + block.instructions = new_instructions; + } + } + fn add_checks_for_loads_of_uninitialized_variables(&mut self) { let nlocals = self.metadata.varnames.len(); if nlocals == 0 { return; } + let merged_cell_local = |cell_relative: usize| { + self.metadata + .cellvars + .get_index(cell_relative) + .and_then(|name| self.metadata.varnames.get_index_of(name.as_str())) + }; + let mut nparams = self.metadata.argcount as usize + self.metadata.kwonlyargcount as usize; if self.flags.contains(CodeFlags::VARARGS) { nparams += 1; @@ -1967,6 +3168,12 @@ impl CodeInfo { worklist.push(target); } } + if matches!(info.instr.real(), Some(Instruction::ForIter { .. })) + && info.target != BlockIdx::NULL + && merge_unsafe_mask(&mut in_masks[info.target.idx()], &unsafe_mask) + { + worklist.push(info.target); + } match info.instr.real() { Some(Instruction::DeleteFast { var_num }) => { let var_idx = usize::from(var_num.get(info.arg)); @@ -1989,6 +3196,15 @@ impl CodeInfo { } new_instructions.push(info); } + Some(Instruction::StoreDeref { i }) => { + let cell_relative = usize::from(i.get(info.arg)); + if let Some(var_idx) = merged_cell_local(cell_relative) + && var_idx < nlocals + { + unsafe_mask[var_idx] = false; + } + new_instructions.push(info); + } Some(Instruction::StoreFastStoreFast { var_nums }) => { let packed = var_nums.get(info.arg); let (idx1, idx2) = packed.indexes(); @@ -2009,7 +3225,17 @@ impl CodeInfo { } new_instructions.push(info); } - Some(Instruction::LoadFast { var_num }) => { + Some(Instruction::DeleteDeref { i }) => { + let cell_relative = usize::from(i.get(info.arg)); + if let Some(var_idx) = merged_cell_local(cell_relative) + && var_idx < nlocals + { + unsafe_mask[var_idx] = true; + } + new_instructions.push(info); + } + Some(Instruction::LoadFast { var_num }) + | Some(Instruction::LoadFastBorrow { var_num }) => { let var_idx = usize::from(var_num.get(info.arg)); if var_idx < nlocals && unsafe_mask[var_idx] { info.instr = Opcode::LoadFastCheck.into(); @@ -2020,7 +3246,8 @@ impl CodeInfo { } new_instructions.push(info); } - Some(Instruction::LoadFastLoadFast { var_nums }) => { + Some(Instruction::LoadFastLoadFast { var_nums }) + | Some(Instruction::LoadFastBorrowLoadFastBorrow { var_nums }) => { let packed = var_nums.get(info.arg); let (idx1, idx2) = packed.indexes(); let idx1 = usize::from(idx1); @@ -2140,7 +3367,10 @@ impl CodeInfo { if target_depth > maxdepth { maxdepth = target_depth; } - stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth); + let target = next_nonempty_block(&self.blocks, ins.target); + if target != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, target, target_depth); + } } depth = new_depth; if instr.is_scope_exit() || instr.is_unconditional_jump() { @@ -2148,37 +3378,293 @@ impl CodeInfo { } } // Only push next block if it's not NULL - if block.next != BlockIdx::NULL { - stackdepth_push(&mut stack, &mut start_depths, block.next, depth); + let next = next_nonempty_block(&self.blocks, block.next); + if next != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, next, depth); + } + } + if DEBUG { + eprintln!("DONE: {maxdepth}"); + } + + for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) { + block.start_depth = (start_depth != u32::MAX).then_some(start_depth); + } + + // Fix up handler stack_depth in ExceptHandlerInfo using start_depths + // computed above: depth = start_depth - 1 - preserve_lasti + for block in self.blocks.iter_mut() { + for ins in &mut block.instructions { + if let Some(ref mut handler) = ins.except_handler { + let h_start = start_depths[handler.handler_block.idx()]; + if h_start != u32::MAX { + let adjustment = 1 + handler.preserve_lasti as u32; + debug_assert!( + h_start >= adjustment, + "handler start depth {h_start} too shallow for adjustment {adjustment}" + ); + handler.stack_depth = h_start.saturating_sub(adjustment); + } + } } } - if DEBUG { - eprintln!("DONE: {maxdepth}"); + + Ok(maxdepth) + } +} + +#[cfg(test)] +impl CodeInfo { + fn debug_block_dump(&self) -> String { + let mut out = String::new(); + for (block_idx, block) in iter_blocks(&self.blocks) { + use core::fmt::Write; + let _ = writeln!( + out, + "block {} next={} cold={} except={} preserve_lasti={} disable_borrow={} start_depth={}", + u32::from(block_idx), + if block.next == BlockIdx::NULL { + String::from("NULL") + } else { + u32::from(block.next).to_string() + }, + block.cold, + block.except_handler, + block.preserve_lasti, + block.disable_load_fast_borrow, + block + .start_depth + .map(|depth| depth.to_string()) + .unwrap_or_else(|| String::from("None")), + ); + for info in &block.instructions { + let lineno = instruction_lineno(info); + let _ = writeln!( + out, + " [disp={} raw={} override={:?}] {:?} arg={} target={}", + lineno, + info.location.line.get(), + info.lineno_override, + info.instr, + u32::from(info.arg), + if info.target == BlockIdx::NULL { + String::from("NULL") + } else { + u32::from(info.target).to_string() + } + ); + } + } + out + } + + pub(crate) fn debug_late_cfg_trace(mut self) -> crate::InternalResult> { + let mut trace = Vec::new(); + trace.push(("initial".to_owned(), self.debug_block_dump())); + + self.splice_annotations_blocks(); + self.fold_binop_constants(); + self.fold_unary_constants(); + self.fold_binop_constants(); + self.fold_tuple_constants(); + self.fold_list_constants(); + self.fold_set_constants(); + self.fold_const_iterable_for_iter(); + self.convert_to_load_small_int(); + self.remove_unused_consts(); + self.dce(); + self.optimize_build_tuple_unpack(); + self.eliminate_dead_stores(); + self.apply_static_swaps(); + self.peephole_optimize(); + trace.push(( + "after_peephole_optimize".to_owned(), + self.debug_block_dump(), + )); + split_blocks_at_jumps(&mut self.blocks); + trace.push(( + "after_split_blocks_at_jumps".to_owned(), + self.debug_block_dump(), + )); + mark_except_handlers(&mut self.blocks); + label_exception_targets(&mut self.blocks); + jump_threading(&mut self.blocks); + trace.push(("after_jump_threading".to_owned(), self.debug_block_dump())); + self.eliminate_unreachable_blocks(); + self.remove_nops(); + trace.push(( + "after_early_remove_nops".to_owned(), + self.debug_block_dump(), + )); + self.add_checks_for_loads_of_uninitialized_variables(); + self.insert_superinstructions(); + push_cold_blocks_to_end(&mut self.blocks); + + trace.push(( + "after_push_cold_blocks_to_end".to_owned(), + self.debug_block_dump(), + )); + + normalize_jumps(&mut self.blocks); + trace.push(("after_normalize_jumps".to_owned(), self.debug_block_dump())); + + reorder_conditional_exit_and_jump_blocks(&mut self.blocks); + reorder_conditional_jump_and_exit_blocks(&mut self.blocks); + reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); + trace.push(("after_reorder".to_owned(), self.debug_block_dump())); + + inline_small_or_no_lineno_blocks(&mut self.blocks); + trace.push(( + "after_inline_small_or_no_lineno_blocks".to_owned(), + self.debug_block_dump(), + )); + + self.dce(); + self.eliminate_unreachable_blocks(); + trace.push(("after_dce_unreachable".to_owned(), self.debug_block_dump())); + + resolve_line_numbers(&mut self.blocks); + trace.push(( + "after_resolve_line_numbers".to_owned(), + self.debug_block_dump(), + )); + + redirect_empty_block_targets(&mut self.blocks); + trace.push(( + "after_redirect_empty_block_targets".to_owned(), + self.debug_block_dump(), + )); + + duplicate_end_returns(&mut self.blocks, &self.metadata); + trace.push(( + "after_duplicate_end_returns".to_owned(), + self.debug_block_dump(), + )); + + self.dce(); + self.eliminate_unreachable_blocks(); + trace.push(( + "after_second_dce_unreachable".to_owned(), + self.debug_block_dump(), + )); + + remove_redundant_nops_and_jumps(&mut self.blocks); + trace.push(( + "after_remove_redundant_nops_and_jumps".to_owned(), + self.debug_block_dump(), + )); + + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + mark_except_handlers(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); + let _ = self.max_stackdepth()?; + convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); + trace.push(( + "after_convert_pseudo_ops".to_owned(), + self.debug_block_dump(), + )); + self.compute_load_fast_start_depths(); + trace.push(( + "after_compute_load_fast_start_depths".to_owned(), + self.debug_block_dump(), + )); + self.optimize_load_fast_borrow(); + trace.push(( + "after_optimize_load_fast_borrow".to_owned(), + self.debug_block_dump(), + )); + self.deoptimize_borrow_after_push_exc_info(); + self.deoptimize_borrow_for_handler_return_paths(); + self.deoptimize_borrow_for_match_keys_attr(); + trace.push(("after_borrow_deopts".to_owned(), self.debug_block_dump())); + + Ok(trace) + } +} + +impl CodeInfo { + fn remap_block_idx(idx: BlockIdx, base: u32) -> BlockIdx { + if idx == BlockIdx::NULL { + idx + } else { + BlockIdx::new(u32::from(idx) + base) + } + } + + fn splice_annotations_blocks(&mut self) { + let mut placeholder = None; + for (block_idx, block) in self.blocks.iter().enumerate() { + if let Some(instr_idx) = block.instructions.iter().position(|info| { + matches!( + info.instr.pseudo(), + Some(PseudoInstruction::AnnotationsPlaceholder) + ) + }) { + placeholder = Some((block_idx, instr_idx)); + break; + } } - for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) { - block.start_depth = (start_depth != u32::MAX).then_some(start_depth); + let Some((block_idx, instr_idx)) = placeholder else { + return; + }; + + let Some(mut annotations_blocks) = self.annotations_blocks.take() else { + self.blocks[block_idx].instructions.remove(instr_idx); + return; + }; + if annotations_blocks.is_empty() { + self.blocks[block_idx].instructions.remove(instr_idx); + return; } - // Fix up handler stack_depth in ExceptHandlerInfo using start_depths - // computed above: depth = start_depth - 1 - preserve_lasti - for block in self.blocks.iter_mut() { - for ins in &mut block.instructions { - if let Some(ref mut handler) = ins.except_handler { - let h_start = start_depths[handler.handler_block.idx()]; - if h_start != u32::MAX { - let adjustment = 1 + handler.preserve_lasti as u32; - debug_assert!( - h_start >= adjustment, - "handler start depth {h_start} too shallow for adjustment {adjustment}" - ); - handler.stack_depth = h_start.saturating_sub(adjustment); - } + let base = self.blocks.len() as u32; + for block in &mut annotations_blocks { + block.next = Self::remap_block_idx(block.next, base); + for info in &mut block.instructions { + info.target = Self::remap_block_idx(info.target, base); + if let Some(handler) = &mut info.except_handler { + handler.handler_block = Self::remap_block_idx(handler.handler_block, base); } } } - Ok(maxdepth) + let ann_entry = BlockIdx::new(base); + let ann_tail = { + let mut cursor = ann_entry; + while annotations_blocks[(u32::from(cursor) - base) as usize].next != BlockIdx::NULL { + cursor = annotations_blocks[(u32::from(cursor) - base) as usize].next; + } + cursor + }; + + let old_next = self.blocks[block_idx].next; + let suffix = self.blocks[block_idx].instructions.split_off(instr_idx + 1); + self.blocks[block_idx].instructions.pop(); + + let suffix_block = if suffix.is_empty() { + old_next + } else { + let suffix_idx = BlockIdx::new(base + annotations_blocks.len() as u32); + let disable_load_fast_borrow = self.blocks[block_idx].disable_load_fast_borrow; + let block = Block { + instructions: suffix, + next: old_next, + disable_load_fast_borrow, + ..Default::default() + }; + annotations_blocks.push(block); + suffix_idx + }; + + self.blocks[block_idx].next = ann_entry; + let ann_tail_local = (u32::from(ann_tail) - base) as usize; + annotations_blocks[ann_tail_local].next = suffix_block; + self.blocks.extend(annotations_blocks); } } @@ -2543,6 +4029,7 @@ fn push_cold_blocks_to_end(blocks: &mut Vec) { location: SourceLocation::default(), end_location: SourceLocation::default(), except_handler: None, + folded_from_nonliteral_expr: false, lineno_override: Some(-1), cache_entries: 0, }); @@ -2624,11 +4111,13 @@ fn split_blocks_at_jumps(blocks: &mut Vec) { let tail: Vec = blocks[bi].instructions.drain(pos..).collect(); let old_next = blocks[bi].next; let cold = blocks[bi].cold; + let disable_load_fast_borrow = blocks[bi].disable_load_fast_borrow; blocks[bi].next = new_block_idx; blocks.push(Block { instructions: tail, next: old_next, cold, + disable_load_fast_borrow, ..Block::default() }); // Don't increment bi - re-check current block (it might still have issues) @@ -2676,14 +4165,28 @@ fn threaded_jump_instr( } let source_kind = jump_thread_kind(source)?; - if source_kind == JumpThreadKind::NoInterrupt { - return Some(source); - } + let result_kind = if source_kind == JumpThreadKind::NoInterrupt + && target_kind == JumpThreadKind::NoInterrupt + { + JumpThreadKind::NoInterrupt + } else { + JumpThreadKind::Plain + }; - Some(match source.into() { - AnyOpcode::Pseudo(_) => PseudoOpcode::Jump.into(), - AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt) => Opcode::JumpBackward.into(), - AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward) => source, + Some(match (source.into(), result_kind) { + (AnyOpcode::Pseudo(_), JumpThreadKind::Plain) => PseudoOpcode::Jump.into(), + (AnyOpcode::Pseudo(_), JumpThreadKind::NoInterrupt) => PseudoOpcode::JumpNoInterrupt.into(), + (AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt), JumpThreadKind::Plain) => { + Opcode::JumpBackward.into() + } + (AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt), JumpThreadKind::NoInterrupt) => source, + (AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward), JumpThreadKind::Plain) => { + source + } + ( + AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward), + JumpThreadKind::NoInterrupt, + ) => PseudoOpcode::JumpNoInterrupt.into(), _ => return None, }) } @@ -2722,11 +4225,13 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { if target == BlockIdx::NULL { continue; } - // Check if target block's first instruction is an unconditional jump + // Thread through blocks that are only leading NOPs followed by an + // unconditional jump so late line anchors do not leave + // JUMP_FORWARD -> NOP -> JUMP_BACKWARD chains behind. let target_jump = blocks[target.idx()] .instructions .iter() - .find(|ins| !matches!(ins.instr.real(), Some(Instruction::Nop))) + .find(|info| !matches!(info.instr.real(), Some(Instruction::Nop))) .copied(); if let Some(target_ins) = target_jump && target_ins.instr.is_unconditional_jump() @@ -2830,6 +4335,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: last_ins.location, end_location: last_ins.end_location, except_handler: last_ins.except_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }; @@ -2845,11 +4351,13 @@ fn normalize_jumps(blocks: &mut Vec) { if let Some(reversed) = reversed_conditional(&last_ins.instr) { let old_next = blocks[idx].next; let is_cold = blocks[idx].cold; + let disable_load_fast_borrow = blocks[idx].disable_load_fast_borrow; // Create new block with NOT_TAKEN + JUMP to original backward target let new_block_idx = BlockIdx(blocks.len() as u32); let mut new_block = Block { cold: is_cold, + disable_load_fast_borrow, ..Block::default() }; new_block.instructions.push(InstructionInfo { @@ -2859,6 +4367,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: loc, end_location: end_loc, except_handler: exc_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); @@ -2869,6 +4378,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: loc, end_location: end_loc, except_handler: exc_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); @@ -2949,7 +4459,30 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { .iter() .all(|ins| !instruction_has_lineno(ins)) }; - + let current_is_named_except_cleanup_normal_exit = |block: &Block| { + let len = block.instructions.len(); + if len < 5 { + return false; + } + let tail = &block.instructions[len - 5..]; + matches!(tail[0].instr.real(), Some(Instruction::PopExcept)) + && matches!(tail[1].instr.real(), Some(Instruction::LoadConst { .. })) + && matches!( + tail[2].instr.real(), + Some(Instruction::StoreName { .. } | Instruction::StoreFast { .. }) + ) + && matches!( + tail[3].instr.real(), + Some(Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }) + ) + && tail[4].instr.is_unconditional_jump() + }; + let target_pushes_handler = |block: &Block| { + block + .instructions + .iter() + .any(|ins| ins.instr.is_block_push()) + }; loop { let mut changes = false; let mut current = BlockIdx(0); @@ -2967,6 +4500,8 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { let target = last.target; if block_is_exceptional(&blocks[current.idx()]) || block_is_exceptional(&blocks[target.idx()]) + || (current_is_named_except_cleanup_normal_exit(&blocks[current.idx()]) + && target_pushes_handler(&blocks[target.idx()])) { current = next; continue; @@ -2975,12 +4510,16 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { && blocks[target.idx()].instructions.len() <= MAX_COPY_SIZE; let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()]) && !block_has_fallthrough(&blocks[target.idx()]); - if small_exit_block || no_lineno_no_fallthrough { + let removed_jump_location = last.location; + let removed_jump_end_location = last.end_location; if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { set_to_nop(last_instr); } - let appended = blocks[target.idx()].instructions.clone(); + let mut appended = blocks[target.idx()].instructions.clone(); + if let Some(first) = appended.first_mut() { + overwrite_location(first, removed_jump_location, removed_jump_end_location); + } blocks[current.idx()].instructions.extend(appended); changes = true; } @@ -3018,28 +4557,43 @@ fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { if lineno < 0 || prev_lineno == lineno { remove = true; } else if src < src_instructions.len() - 1 { - let next_lineno = instruction_lineno(&src_instructions[src + 1]); - if next_lineno == lineno { - remove = true; - } else if next_lineno < 0 { + if src_instructions[src + 1].instr.is_unconditional_jump() { src_instructions[src + 1].lineno_override = Some(lineno); remove = true; + } else if src_instructions[src + 1].folded_from_nonliteral_expr { + remove = true; + } else { + let next_lineno = instruction_lineno(&src_instructions[src + 1]); + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } } } else { let next = next_nonempty_block(blocks, blocks[bi].next); if next != BlockIdx::NULL { - let mut next_lineno = None; - for next_instr in &blocks[next.idx()].instructions { + let mut next_info = None; + for (next_idx, next_instr) in + blocks[next.idx()].instructions.iter().enumerate() + { let line = instruction_lineno(next_instr); if matches!(next_instr.instr.real(), Some(Instruction::Nop)) && line < 0 { continue; } - next_lineno = Some(line); + next_info = Some((next_idx, line)); break; } - if next_lineno.is_some_and(|line| line == lineno) { - remove = true; + if let Some((next_idx, next_lineno)) = next_info { + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + blocks[next.idx()].instructions[next_idx].lineno_override = + Some(lineno); + remove = true; + } } } } @@ -3093,6 +4647,33 @@ fn remove_redundant_nops_and_jumps(blocks: &mut [Block]) { } } +fn redirect_empty_block_targets(blocks: &mut [Block]) { + let redirected_targets: Vec> = blocks + .iter() + .map(|block| { + block + .instructions + .iter() + .map(|instr| { + if instr.target == BlockIdx::NULL { + BlockIdx::NULL + } else { + next_nonempty_block(blocks, instr.target) + } + }) + .collect() + }) + .collect(); + + for (block, block_targets) in blocks.iter_mut().zip(redirected_targets) { + for (instr, target) in block.instructions.iter_mut().zip(block_targets) { + if target != BlockIdx::NULL { + instr.target = target; + } + } + } +} + fn merge_unsafe_mask(slot: &mut Option>, incoming: &[bool]) -> bool { match slot { Some(existing) => { @@ -3112,6 +4693,38 @@ fn merge_unsafe_mask(slot: &mut Option>, incoming: &[bool]) -> bool { } } +fn deoptimize_borrow_after_push_exc_info_in_blocks(blocks: &mut [Block]) { + let mut in_exception_state = false; + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let block = &mut blocks[current.idx()]; + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::PushExcInfo) => { + in_exception_state = true; + } + Some(Instruction::PopExcept) | Some(Instruction::Reraise { .. }) => { + in_exception_state = false; + } + Some(Instruction::LoadFastBorrow { .. }) if in_exception_state => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) if in_exception_state => { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + current = block.next; + } +} + /// Follow chain of empty blocks to find first non-empty block. fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx { while idx != BlockIdx::NULL @@ -3123,6 +4736,14 @@ fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx { idx } +fn is_load_const_none(instr: &InstructionInfo, metadata: &CodeUnitMetadata) -> bool { + matches!(instr.instr.real(), Some(Instruction::LoadConst { .. })) + && matches!( + metadata.consts.get_index(u32::from(instr.arg) as usize), + Some(ConstantData::None) + ) +} + fn instruction_lineno(instr: &InstructionInfo) -> i32 { instr .lineno_override @@ -3447,6 +5068,7 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { let mut target_end = BlockIdx::NULL; let mut target_exit = BlockIdx::NULL; + let mut nonempty_target_blocks = 0usize; cursor = target; while cursor != BlockIdx::NULL { if block_is_exceptional(&blocks[cursor.idx()]) { @@ -3454,6 +5076,7 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { } target_end = cursor; if !blocks[cursor.idx()].instructions.is_empty() { + nonempty_target_blocks += 1; target_exit = cursor; } cursor = blocks[cursor.idx()].next; @@ -3461,6 +5084,8 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { if target_end == BlockIdx::NULL || target_exit == BlockIdx::NULL + || nonempty_target_blocks != 1 + || target_exit != target_end || !is_scope_exit_block(&blocks[target_exit.idx()]) { current = next; @@ -3487,6 +5112,16 @@ fn maybe_propagate_location( } } +fn overwrite_location( + instr: &mut InstructionInfo, + location: SourceLocation, + end_location: SourceLocation, +) { + instr.location = location; + instr.end_location = end_location; + instr.lineno_override = None; +} + fn propagate_locations_in_block( block: &mut Block, location: SourceLocation, @@ -3706,27 +5341,37 @@ fn find_layout_predecessor(blocks: &[Block], target: BlockIdx) -> BlockIdx { /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through /// to the final return block. -fn duplicate_end_returns(blocks: &mut Vec) { - // Walk the block chain and keep the last non-empty block. +fn duplicate_end_returns(blocks: &mut Vec, metadata: &CodeUnitMetadata) { + // Walk the block chain and keep the last non-cold non-empty block. + // After cold exception handlers are pushed to the end, the mainline + // return epilogue can sit before trailing cold blocks. let mut last_block = BlockIdx::NULL; + let mut last_nonempty_block = BlockIdx::NULL; let mut current = BlockIdx(0); while current != BlockIdx::NULL { if !blocks[current.idx()].instructions.is_empty() { - last_block = current; + last_nonempty_block = current; + if !blocks[current.idx()].cold { + last_block = current; + } } current = blocks[current.idx()].next; } + if last_block == BlockIdx::NULL { + last_block = last_nonempty_block; + } if last_block == BlockIdx::NULL { return; } let last_insts = &blocks[last_block.idx()].instructions; - // Only apply when the last block is EXACTLY a return-None epilogue + // Only apply when the last block is EXACTLY a return-None epilogue. let is_return_block = last_insts.len() == 2 && matches!( last_insts[0].instr, AnyInstruction::Real(Instruction::LoadConst { .. }) ) + && is_load_const_none(&last_insts[0], metadata) && matches!( last_insts[1].instr, AnyInstruction::Real(Instruction::ReturnValue) @@ -3748,7 +5393,7 @@ fn duplicate_end_returns(blocks: &mut Vec) { while current != BlockIdx::NULL { let block = &blocks[current.idx()]; let next = next_nonempty_block(blocks, block.next); - if current != last_block && !block.cold && !block.except_handler { + if current != last_block && !block.cold { let last_ins = block.instructions.last(); let has_fallthrough = last_ins .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()) @@ -3764,20 +5409,28 @@ fn duplicate_end_returns(blocks: &mut Vec) { AnyInstruction::Real(Instruction::ReturnValue) ) }; - if next == last_block + if !block.except_handler + && next == last_block && has_fallthrough && trailing_conditional_jump_index(block).is_none() && !already_has_return { fallthrough_blocks_to_fix.push(current); } - if predecessors[last_block.idx()] > 1 - && let Some(last) = block.instructions.last() - && last.instr.is_unconditional_jump() - && last.target != BlockIdx::NULL - && next_nonempty_block(blocks, last.target) == last_block - { - jump_targets_to_fix.push((current, block.instructions.len() - 1)); + let jump_idx = trailing_conditional_jump_index(block).or_else(|| { + block.instructions.last().and_then(|last| { + (last.instr.is_unconditional_jump() && last.target != BlockIdx::NULL) + .then_some(block.instructions.len() - 1) + }) + }); + if let Some(jump_idx) = jump_idx { + let jump = &block.instructions[jump_idx]; + if jump.target != BlockIdx::NULL + && next_nonempty_block(blocks, jump.target) == last_block + && (is_conditional_jump(&jump.instr) || predecessors[last_block.idx()] > 1) + { + jump_targets_to_fix.push((current, jump_idx)); + } } } current = blocks[current.idx()].next; @@ -3785,24 +5438,33 @@ fn duplicate_end_returns(blocks: &mut Vec) { // Duplicate the return instructions at the end of fall-through blocks for block_idx in fallthrough_blocks_to_fix { - blocks[block_idx.idx()] + let propagated_location = blocks[block_idx.idx()] .instructions - .extend_from_slice(&return_insts); + .last() + .map(|instr| (instr.location, instr.end_location)); + let mut cloned_return = return_insts.clone(); + if let Some((location, end_location)) = propagated_location { + for instr in &mut cloned_return { + overwrite_location(instr, location, end_location); + } + } + blocks[block_idx.idx()].instructions.extend(cloned_return); } // Clone the final return block for jump predecessors so their target layout // matches CPython's duplicated exit blocks. - for (block_idx, instr_idx) in jump_targets_to_fix { + for (block_idx, instr_idx) in jump_targets_to_fix.into_iter().rev() { let jump = blocks[block_idx.idx()].instructions[instr_idx]; let mut cloned_return = return_insts.clone(); - for instr in &mut cloned_return { - maybe_propagate_location(instr, jump.location, jump.end_location); + if let Some(first) = cloned_return.first_mut() { + overwrite_location(first, jump.location, jump.end_location); } let new_idx = BlockIdx(blocks.len() as u32); let is_conditional = is_conditional_jump(&jump.instr); let new_block = Block { cold: blocks[last_block.idx()].cold, except_handler: blocks[last_block.idx()].except_handler, + disable_load_fast_borrow: blocks[last_block.idx()].disable_load_fast_borrow, instructions: cloned_return, next: if is_conditional { last_block @@ -3824,6 +5486,182 @@ fn duplicate_end_returns(blocks: &mut Vec) { } } +fn inline_with_suppress_return_blocks(blocks: &mut [Block]) { + fn has_with_suppress_prefix(block: &Block, jump_idx: usize) -> bool { + let tail: Vec<_> = block.instructions[..jump_idx] + .iter() + .filter_map(|info| info.instr.real()) + .rev() + .take(5) + .collect(); + matches!( + tail.as_slice(), + [ + Instruction::PopTop, + Instruction::PopTop, + Instruction::PopTop, + Instruction::PopExcept, + Instruction::PopTop, + ] + ) + } + + for block_idx in 0..blocks.len() { + let Some(jump_idx) = blocks[block_idx].instructions.len().checked_sub(1) else { + continue; + }; + let jump = blocks[block_idx].instructions[jump_idx]; + if !jump.instr.is_unconditional_jump() || jump.target == BlockIdx::NULL { + continue; + } + if !has_with_suppress_prefix(&blocks[block_idx], jump_idx) { + continue; + } + + let target = next_nonempty_block(blocks, jump.target); + if target == BlockIdx::NULL || !is_const_return_block(&blocks[target.idx()]) { + continue; + } + + let mut cloned_return = blocks[target.idx()].instructions.clone(); + for instr in &mut cloned_return { + overwrite_location(instr, jump.location, jump.end_location); + } + blocks[block_idx].instructions.pop(); + blocks[block_idx].instructions.extend(cloned_return); + } +} + +fn is_named_except_cleanup_return_block(block: &Block, metadata: &CodeUnitMetadata) -> bool { + matches!( + block.instructions.as_slice(), + [pop_except, load_none1, store, delete, load_none2, ret] + if matches!(pop_except.instr.real(), Some(Instruction::PopExcept)) + && is_load_const_none(load_none1, metadata) + && matches!( + store.instr.real(), + Some(Instruction::StoreFast { .. } | Instruction::StoreName { .. }) + ) + && matches!( + delete.instr.real(), + Some(Instruction::DeleteFast { .. } | Instruction::DeleteName { .. }) + ) + && is_load_const_none(load_none2, metadata) + && matches!(ret.instr.real(), Some(Instruction::ReturnValue)) + ) +} + +fn duplicate_named_except_cleanup_returns(blocks: &mut Vec, metadata: &CodeUnitMetadata) { + let predecessors = compute_predecessors(blocks); + let mut clones = Vec::new(); + + for target in 0..blocks.len() { + let target = BlockIdx(target as u32); + if !is_named_except_cleanup_return_block(&blocks[target.idx()], metadata) { + continue; + } + + let layout_pred = find_layout_predecessor(blocks, target); + if layout_pred == BlockIdx::NULL + || next_nonempty_block(blocks, blocks[layout_pred.idx()].next) != target + { + continue; + } + + let fallthroughs_into_target = blocks[layout_pred.idx()] + .instructions + .last() + .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()) + .unwrap_or(true); + if !fallthroughs_into_target || predecessors[target.idx()] < 2 { + continue; + } + + for block_idx in 0..blocks.len() { + if block_idx == target.idx() { + continue; + } + let Some(instr_idx) = trailing_conditional_jump_index(&blocks[block_idx]) else { + continue; + }; + if next_nonempty_block(blocks, blocks[block_idx].instructions[instr_idx].target) + != target + { + continue; + } + clones.push((BlockIdx(block_idx as u32), instr_idx, target)); + } + } + + for (block_idx, instr_idx, target) in clones.into_iter().rev() { + let jump = blocks[block_idx.idx()].instructions[instr_idx]; + let mut cloned = blocks[target.idx()].instructions.clone(); + if let Some(first) = cloned.first_mut() { + overwrite_location(first, jump.location, jump.end_location); + } + + let new_idx = BlockIdx(blocks.len() as u32); + let next = blocks[target.idx()].next; + blocks.push(Block { + cold: blocks[target.idx()].cold, + except_handler: blocks[target.idx()].except_handler, + disable_load_fast_borrow: blocks[target.idx()].disable_load_fast_borrow, + instructions: cloned, + next, + ..Block::default() + }); + blocks[target.idx()].next = new_idx; + blocks[block_idx.idx()].instructions[instr_idx].target = new_idx; + } +} + +fn is_const_return_block(block: &Block) -> bool { + block.instructions.len() == 2 + && matches!( + block.instructions[0].instr.real(), + Some(Instruction::LoadConst { .. }) + ) + && matches!( + block.instructions[1].instr.real(), + Some(Instruction::ReturnValue) + ) +} + +fn inline_pop_except_return_blocks(blocks: &mut [Block]) { + for block_idx in 0..blocks.len() { + let Some(jump_idx) = blocks[block_idx].instructions.len().checked_sub(1) else { + continue; + }; + let jump = blocks[block_idx].instructions[jump_idx]; + if !jump.instr.is_unconditional_jump() || jump.target == BlockIdx::NULL { + continue; + } + + let Some(last_real_before_jump) = blocks[block_idx].instructions[..jump_idx] + .iter() + .rev() + .find_map(|info| info.instr.real()) + else { + continue; + }; + if !matches!(last_real_before_jump, Instruction::PopExcept) { + continue; + } + + let target = next_nonempty_block(blocks, jump.target); + if target == BlockIdx::NULL || !is_const_return_block(&blocks[target.idx()]) { + continue; + } + + let mut cloned_return = blocks[target.idx()].instructions.clone(); + for instr in &mut cloned_return { + overwrite_location(instr, jump.location, jump.end_location); + } + blocks[block_idx].instructions.pop(); + blocks[block_idx].instructions.extend(cloned_return); + } +} + /// Label exception targets: walk CFG with except stack, set per-instruction /// handler info and block preserve_lasti flag. Converts POP_BLOCK to NOP. /// flowgraph.c label_exception_targets + push_except_block diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap new file mode 100644 index 00000000000..840f4397d75 --- /dev/null +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap @@ -0,0 +1,67 @@ +--- +source: crates/codegen/src/compile.rs +assertion_line: 12138 +expression: "compile_exec(\"\\\ndef f(one: int):\n int.new_attr: int\n [list][0].new_attr: [int, str]\n my_lst = [1]\n my_lst[one]: int\n return my_lst\n\")" +--- + 1 0 RESUME (0) + 1 LOAD_CONST (): 1 0 RESUME (0) + 1 LOAD_FAST_BORROW (0, format) + 2 LOAD_SMALL_INT (2) + >> 3 COMPARE_OP (>) + 4 CACHE + 5 POP_JUMP_IF_FALSE (3) + 6 CACHE + 7 NOT_TAKEN + 8 LOAD_COMMON_CONSTANT (NotImplementedError) + 9 RAISE_VARARGS (Raise) + 10 LOAD_CONST ("one") + 11 LOAD_GLOBAL (0, int) + 12 CACHE + 13 CACHE + 14 CACHE + 15 CACHE + 16 BUILD_MAP (1) + 17 RETURN_VALUE + + 2 MAKE_FUNCTION + 3 LOAD_CONST (): 1 0 RESUME (0) + + 2 1 LOAD_GLOBAL (0, int) + 2 CACHE + 3 CACHE + 4 CACHE + 5 CACHE + 6 POP_TOP + + 3 7 LOAD_GLOBAL (2, list) + 8 CACHE + 9 CACHE + 10 CACHE + 11 CACHE + 12 BUILD_LIST (1) + 13 LOAD_SMALL_INT (0) + 14 BINARY_OP ([]) + 15 CACHE + 16 CACHE + 17 CACHE + 18 CACHE + 19 CACHE + 20 POP_TOP + + 4 21 LOAD_SMALL_INT (1) + 22 BUILD_LIST (1) + 23 STORE_FAST (1, my_lst) + + 5 24 LOAD_FAST_BORROW (1, my_lst) + 25 POP_TOP + 26 LOAD_FAST_BORROW (0, one) + 27 POP_TOP + + 6 28 LOAD_FAST_BORROW (1, my_lst) + 29 RETURN_VALUE + + 4 MAKE_FUNCTION + 5 SET_FUNCTION_ATTRIBUTE(Annotate) + 6 STORE_NAME (0, f) + 7 LOAD_CONST (None) + 8 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap new file mode 100644 index 00000000000..a600a829863 --- /dev/null +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap @@ -0,0 +1,10 @@ +--- +source: crates/codegen/src/compile.rs +assertion_line: 12222 +expression: "compile_exec(\"\\\nif 1:\n pass\n\")" +--- + 1 0 RESUME (0) + 1 NOP + + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap index 0cdb7f0a3df..8e9bb5d25f4 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap @@ -1,23 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10890 +assertion_line: 11413 expression: "compile_exec(\"\\\nif True and False and False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (11) - >> 3 CACHE - 4 NOT_TAKEN - 5 LOAD_CONST (False) - 6 POP_JUMP_IF_FALSE (7) - >> 7 CACHE - 8 NOT_TAKEN - 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - >> 11 CACHE - 12 NOT_TAKEN - - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 RETURN_VALUE + 1 LOAD_CONST (None) + 2 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap index 8e91189912d..f7df6f4f3ee 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap @@ -1,27 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10900 +assertion_line: 11423 expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (5) - >> 3 CACHE - 4 NOT_TAKEN - >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (9) - >> 7 CACHE - 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (7) - 11 CACHE - 12 NOT_TAKEN - 13 LOAD_CONST (True) - 14 POP_JUMP_IF_FALSE (3) - 15 CACHE - 16 NOT_TAKEN - - 2 17 LOAD_CONST (None) - 18 RETURN_VALUE - 19 LOAD_CONST (None) - 20 RETURN_VALUE + 1 LOAD_CONST (None) + 2 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap index 9445635458d..f38d3c2c593 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap @@ -1,23 +1,10 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10880 +assertion_line: 11039 expression: "compile_exec(\"\\\nif True or False or False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_TRUE (9) - >> 3 CACHE - 4 NOT_TAKEN - >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (5) - 7 CACHE - 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - 11 CACHE - 12 NOT_TAKEN + 1 NOP - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 RETURN_VALUE + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap index 64ae0cfd5ff..fccc7b7c336 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap @@ -1,6 +1,6 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10936 +assertion_line: 11626 expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')):\n with self.subTest(type=type(stop_exc)):\n try:\n async with egg():\n raise stop_exc\n except Exception as ex:\n self.assertIs(ex, stop_exc)\n else:\n self.fail(f'{stop_exc} was suppressed')\n\")" --- 1 0 RESUME (0) @@ -13,7 +13,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter >> 5 CACHE 6 CACHE 7 CACHE - 8 LOAD_CONST ("spam") + >> 8 LOAD_CONST ("spam") 9 CALL (1) >> 10 CACHE 11 CACHE @@ -34,7 +34,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 26 CACHE 27 STORE_FAST (0, stop_exc) - 3 >> 28 LOAD_GLOBAL (4, self) + 3 28 LOAD_GLOBAL (4, self) 29 CACHE 30 CACHE 31 CACHE @@ -51,10 +51,10 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 42 CACHE 43 LOAD_GLOBAL (9, NULL + type) 44 CACHE - 45 CACHE + >> 45 CACHE 46 CACHE 47 CACHE - >> 48 LOAD_FAST (0, stop_exc) + 48 LOAD_FAST_BORROW (0, stop_exc) 49 CALL (1) 50 CACHE 51 CACHE @@ -105,7 +105,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 94 END_SEND 95 POP_TOP - 6 96 LOAD_FAST (0, stop_exc) + 6 96 LOAD_FAST_BORROW (0, stop_exc) 97 RAISE_VARARGS (Raise) 2 98 END_FOR @@ -139,115 +139,118 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 125 POP_TOP 126 POP_TOP 127 POP_TOP - 128 JUMP_FORWARD (48) + 128 JUMP_FORWARD (3) 129 COPY (3) 130 POP_EXCEPT 131 RERAISE (1) - 132 PUSH_EXC_INFO + 132 NOP - 7 133 LOAD_GLOBAL (12, Exception) + 10 133 LOAD_GLOBAL (4, self) 134 CACHE 135 CACHE 136 CACHE 137 CACHE - 138 CHECK_EXC_MATCH - 139 POP_JUMP_IF_FALSE (28) + 138 LOAD_ATTR (13, fail, method=true) + 139 CACHE 140 CACHE - 141 NOT_TAKEN - 142 STORE_FAST (1, ex) - - 8 143 LOAD_GLOBAL (4, self) + 141 CACHE + 142 CACHE + 143 CACHE 144 CACHE 145 CACHE 146 CACHE 147 CACHE - 148 LOAD_ATTR (15, assertIs, method=true) - 149 CACHE - 150 CACHE - 151 CACHE - 152 CACHE + 148 LOAD_FAST (0, stop_exc) + 149 FORMAT_SIMPLE + 150 LOAD_CONST (" was suppressed") + 151 BUILD_STRING (2) + 152 CALL (1) 153 CACHE 154 CACHE 155 CACHE - 156 CACHE - 157 CACHE - 158 LOAD_FAST_LOAD_FAST (ex, stop_exc) - 159 CALL (2) + 156 POP_TOP + 157 JUMP_FORWARD (45) + 158 PUSH_EXC_INFO + + 7 159 LOAD_GLOBAL (14, Exception) 160 CACHE 161 CACHE 162 CACHE - 163 POP_TOP - 164 POP_EXCEPT - 165 LOAD_CONST (None) - 166 STORE_FAST (1, ex) - 167 DELETE_FAST (1, ex) - 168 JUMP_FORWARD (32) - 169 RERAISE (0) - 170 LOAD_CONST (None) - 171 STORE_FAST (1, ex) - 172 DELETE_FAST (1, ex) - 173 RERAISE (1) - 174 COPY (3) - 175 POP_EXCEPT - 176 RERAISE (1) + 163 CACHE + 164 CHECK_EXC_MATCH + 165 POP_JUMP_IF_FALSE (32) + 166 CACHE + 167 NOT_TAKEN + 168 STORE_FAST (1, ex) - 10 177 LOAD_GLOBAL (4, self) + 8 169 LOAD_GLOBAL (4, self) + 170 CACHE + 171 CACHE + 172 CACHE + 173 CACHE + 174 LOAD_ATTR (17, assertIs, method=true) + 175 CACHE + 176 CACHE + 177 CACHE 178 CACHE 179 CACHE 180 CACHE 181 CACHE - 182 LOAD_ATTR (17, fail, method=true) + 182 CACHE 183 CACHE - 184 CACHE - 185 CACHE - >> 186 CACHE + 184 LOAD_FAST_LOAD_FAST (ex, stop_exc) + 185 CALL (2) + 186 CACHE 187 CACHE - 188 CACHE - 189 CACHE - 190 CACHE - 191 CACHE - 192 LOAD_FAST_BORROW (0, stop_exc) - 193 FORMAT_SIMPLE - 194 LOAD_CONST (" was suppressed") - 195 BUILD_STRING (2) - 196 CALL (1) - 197 CACHE - 198 CACHE - 199 CACHE - 200 POP_TOP + >> 188 CACHE + 189 POP_TOP + 190 POP_EXCEPT + 191 LOAD_CONST (None) + 192 STORE_FAST (1, ex) + 193 DELETE_FAST (1, ex) + 194 JUMP_FORWARD (8) + 195 LOAD_CONST (None) + 196 STORE_FAST (1, ex) + 197 DELETE_FAST (1, ex) + 198 RERAISE (1) + 199 RERAISE (0) + 200 COPY (3) + 201 POP_EXCEPT + 202 RERAISE (1) - 3 201 LOAD_CONST (None) - >> 202 LOAD_CONST (None) - 203 LOAD_CONST (None) - 204 CALL (3) - 205 CACHE - 206 CACHE + 3 203 LOAD_CONST (None) + 204 LOAD_CONST (None) + >> 205 LOAD_CONST (None) + 206 CALL (3) 207 CACHE - 208 POP_TOP - 209 JUMP_BACKWARD (186) - 210 CACHE - 211 PUSH_EXC_INFO - 212 WITH_EXCEPT_START - 213 TO_BOOL - 214 CACHE - 215 CACHE + 208 CACHE + 209 CACHE + 210 POP_TOP + 211 JUMP_BACKWARD (188) + 212 CACHE + 213 PUSH_EXC_INFO + 214 WITH_EXCEPT_START + 215 TO_BOOL 216 CACHE - 217 POP_JUMP_IF_TRUE (2) + 217 CACHE 218 CACHE - 219 NOT_TAKEN - 220 RERAISE (2) - 221 POP_TOP - 222 POP_EXCEPT + 219 POP_JUMP_IF_TRUE (2) + 220 CACHE + 221 NOT_TAKEN + 222 RERAISE (2) 223 POP_TOP - 224 POP_TOP + 224 POP_EXCEPT 225 POP_TOP - 226 JUMP_BACKWARD_NO_INTERRUPT(202) - 227 COPY (3) - 228 POP_EXCEPT - 229 RERAISE (1) + 226 POP_TOP + 227 POP_TOP + 228 JUMP_BACKWARD (205) + 229 CACHE + 230 COPY (3) + 231 POP_EXCEPT + 232 RERAISE (1) - 2 230 CALL_INTRINSIC_1 (StopIterationError) - 231 RERAISE (1) + 2 233 CALL_INTRINSIC_1 (StopIterationError) + 234 RERAISE (1) 2 MAKE_FUNCTION 3 STORE_NAME (0, test) diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index f8d7a2b26ea..e87aa3d4aea 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -571,10 +571,13 @@ impl SymbolTableAnalyzer { } newfree.extend(child_free); } - if let Some(ann_free) = annotation_free { - // Propagate annotation-scope free names to this scope so - // implicit class-scope cells (__classdict__/__conditional_annotations__) - // can be materialized by drop_class_free when needed. + if let Some(ann_free) = annotation_free + && symbol_table.typ == CompilerScope::Class + { + // Annotation-only free variables should not leak into function + // bodies. We only need to propagate them through class scopes so + // drop_class_free() can materialize implicit class cells when + // annotation scopes reference them. newfree.extend(ann_free); } @@ -1074,6 +1077,9 @@ impl SymbolTableBuilder { // Register .type_params as a SET symbol (it will be converted to cell variable later) self.register_name(".type_params", SymbolUsage::Assigned, TextRange::default())?; + if for_class { + self.register_name(".generic_base", SymbolUsage::Assigned, TextRange::default())?; + } Ok(()) } @@ -1260,27 +1266,38 @@ impl SymbolTableBuilder { is_ann_assign: bool, ) -> SymbolTableResult { let current_scope = self.tables.last().map(|t| t.typ); + let needs_future_annotation_bookkeeping = is_ann_assign + && self.future_annotations + && matches!( + current_scope, + Some(CompilerScope::Module | CompilerScope::Class) + ); + let needs_non_future_conditional_annotations = is_ann_assign + && !self.future_annotations + && (matches!(current_scope, Some(CompilerScope::Module)) + || (matches!(current_scope, Some(CompilerScope::Class)) + && self.in_conditional_block)); + let should_register_conditional_annotations = needs_future_annotation_bookkeeping + || (needs_non_future_conditional_annotations + && !self.tables.last().unwrap().has_conditional_annotations); // PEP 649: Only AnnAssign annotations can be conditional. // Function parameter/return annotations are never conditional. - if is_ann_assign && !self.future_annotations { - let is_conditional = matches!(current_scope, Some(CompilerScope::Module)) - || (matches!(current_scope, Some(CompilerScope::Class)) - && self.in_conditional_block); - - if is_conditional && !self.tables.last().unwrap().has_conditional_annotations { - self.tables.last_mut().unwrap().has_conditional_annotations = true; - self.register_name( - "__conditional_annotations__", - SymbolUsage::Assigned, - annotation.range(), - )?; - self.register_name( - "__conditional_annotations__", - SymbolUsage::Used, - annotation.range(), - )?; - } + if needs_non_future_conditional_annotations { + self.tables.last_mut().unwrap().has_conditional_annotations = true; + } + + if should_register_conditional_annotations { + self.register_name( + "__conditional_annotations__", + SymbolUsage::Assigned, + annotation.range(), + )?; + self.register_name( + "__conditional_annotations__", + SymbolUsage::Used, + annotation.range(), + )?; } // Create annotation scope for deferred evaluation @@ -1434,6 +1451,10 @@ impl SymbolTableBuilder { self.register_name("__qualname__", SymbolUsage::Assigned, *range)?; self.register_name("__doc__", SymbolUsage::Assigned, *range)?; self.register_name("__class__", SymbolUsage::Assigned, *range)?; + if type_params.is_some() { + self.register_name(".type_params", SymbolUsage::Used, *range)?; + self.register_name("__type_params__", SymbolUsage::Assigned, *range)?; + } self.scan_statements(body)?; self.leave_scope(); self.in_conditional_block = saved_in_conditional; @@ -1569,48 +1590,42 @@ impl SymbolTableBuilder { }) => { // https://github.com/python/cpython/blob/main/Python/symtable.c#L1233 match &**target { - Expr::Name(ast::ExprName { id, .. }) if *simple => { + Expr::Name(ast::ExprName { id, .. }) => { let id_str = id.as_str(); - self.check_name(id_str, ExpressionContext::Store, *range)?; - - self.register_name(id_str, SymbolUsage::AnnotationAssigned, *range)?; - // PEP 649: Register annotate function in module/class scope - let current_scope = self.tables.last().map(|t| t.typ); - match current_scope { - Some(CompilerScope::Module) => { - self.register_name("__annotate__", SymbolUsage::Assigned, *range)?; - } - Some(CompilerScope::Class) => { - self.register_name( - "__annotate_func__", - SymbolUsage::Assigned, - *range, - )?; + if *simple { + self.check_name(id_str, ExpressionContext::Store, *range)?; + + self.register_name(id_str, SymbolUsage::AnnotationAssigned, *range)?; + // PEP 649: Register annotate function in module/class scope + let current_scope = self.tables.last().map(|t| t.typ); + match current_scope { + Some(CompilerScope::Module) => { + self.register_name( + "__annotate__", + SymbolUsage::Assigned, + *range, + )?; + } + Some(CompilerScope::Class) => { + self.register_name( + "__annotate_func__", + SymbolUsage::Assigned, + *range, + )?; + } + _ => {} } - _ => {} + } else if value.is_some() { + self.check_name(id_str, ExpressionContext::Store, *range)?; + self.register_name(id_str, SymbolUsage::Assigned, *range)?; } } _ => { self.scan_expression(target, ExpressionContext::Store)?; } } - // Only scan annotation in annotation scope for simple name targets. - // Non-simple annotations (subscript, attribute, parenthesized) are - // never compiled into __annotate__, so scanning them would create - // sub_tables that cause mismatch in the annotation scope's sub_table index. - let is_simple_name = *simple && matches!(&**target, Expr::Name(_)); - if is_simple_name { - self.scan_ann_assign_annotation(annotation)?; - } else { - // Still validate annotation for forbidden expressions - // (yield, await, named) even for non-simple targets. - let was_in_annotation = self.in_annotation; - self.in_annotation = true; - let result = self.scan_expression(annotation, ExpressionContext::Load); - self.in_annotation = was_in_annotation; - result?; - } + self.scan_ann_assign_annotation(annotation)?; if let Some(value) = value { self.scan_expression(value, ExpressionContext::Load)?; } @@ -1639,6 +1654,10 @@ impl SymbolTableBuilder { let saved_in_conditional_block = self.in_conditional_block; self.in_conditional_block = true; self.scan_statements(body)?; + // Preserve source-order symbol analysis so `global`/`nonlocal` + // semantics match CPython, but reorder child scope storage to + // match the codegen order for plain try/except/else. + let body_subtables_len = self.tables.last().unwrap().sub_tables.len(); for handler in handlers { let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { type_, @@ -1654,7 +1673,22 @@ impl SymbolTableBuilder { } self.scan_statements(body)?; } - self.scan_statements(orelse)?; + if finalbody.is_empty() { + let handler_subtables = self + .tables + .last_mut() + .unwrap() + .sub_tables + .split_off(body_subtables_len); + self.scan_statements(orelse)?; + self.tables + .last_mut() + .unwrap() + .sub_tables + .extend(handler_subtables); + } else { + self.scan_statements(orelse)?; + } self.scan_statements(finalbody)?; self.in_conditional_block = saved_in_conditional_block; } @@ -2160,6 +2194,13 @@ impl SymbolTableBuilder { }); } + assert!(!generators.is_empty()); + let outermost = &generators[0]; + + // CPython evaluates the outermost iterator in the enclosing scope + // before entering the comprehension scope. + self.scan_expression(&outermost.iter, ExpressionContext::IterDefinitionExp)?; + // Comprehensions are compiled as functions, so create a scope for them: self.enter_scope( scope_name, @@ -2195,36 +2236,27 @@ impl SymbolTableBuilder { // Register the passed argument to the generator function as the name ".0" self.register_name(".0", SymbolUsage::Parameter, range)?; - self.scan_expression(elt1, ExpressionContext::Load)?; - if let Some(elt2) = elt2 { - self.scan_expression(elt2, ExpressionContext::Load)?; + self.scan_expression(&outermost.target, ExpressionContext::Iter)?; + for if_expr in &outermost.ifs { + self.scan_expression(if_expr, ExpressionContext::Load)?; } - let mut is_first_generator = true; - for generator in generators { - // Set flag for INNER_LOOP_CONFLICT check (only for inner loops, not the first) - if !is_first_generator { - self.in_comp_inner_loop_target = true; - } + for generator in &generators[1..] { + self.in_comp_inner_loop_target = true; self.scan_expression(&generator.target, ExpressionContext::Iter)?; self.in_comp_inner_loop_target = false; - - if is_first_generator { - is_first_generator = false; - } else { - self.scan_expression(&generator.iter, ExpressionContext::IterDefinitionExp)?; - } - + self.scan_expression(&generator.iter, ExpressionContext::IterDefinitionExp)?; for if_expr in &generator.ifs { self.scan_expression(if_expr, ExpressionContext::Load)?; } } - self.leave_scope(); + if let Some(elt2) = elt2 { + self.scan_expression(elt2, ExpressionContext::Load)?; + } + self.scan_expression(elt1, ExpressionContext::Load)?; - // The first iterable is passed as an argument into the created function: - assert!(!generators.is_empty()); - self.scan_expression(&generators[0].iter, ExpressionContext::IterDefinitionExp)?; + self.leave_scope(); Ok(()) } diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index aab176de62a..269fe518e6e 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -614,23 +614,15 @@ impl Instruction { /// Map a specialized opcode back to its adaptive (base) variant. /// `_PyOpcode_Deopt` pub const fn deopt(self) -> Option { - Some(match self { - // RESUME specializations - Self::ResumeCheck => Self::Resume { - context: Arg::marker(), - }, - // LOAD_CONST specializations - Self::LoadConstMortal | Self::LoadConstImmortal => Self::LoadConst { - consti: Arg::marker(), - }, - // TO_BOOL specializations + let opcode = match self { + Self::ResumeCheck => Opcode::Resume, + Self::LoadConstMortal | Self::LoadConstImmortal => Opcode::LoadConst, Self::ToBoolAlwaysTrue | Self::ToBoolBool | Self::ToBoolInt | Self::ToBoolList | Self::ToBoolNone - | Self::ToBoolStr => Self::ToBool, - // BINARY_OP specializations + | Self::ToBoolStr => Opcode::ToBool, Self::BinaryOpMultiplyInt | Self::BinaryOpAddInt | Self::BinaryOpSubtractInt @@ -645,34 +637,18 @@ impl Instruction { | Self::BinaryOpSubscrDict | Self::BinaryOpSubscrGetitem | Self::BinaryOpExtend - | Self::BinaryOpInplaceAddUnicode => Self::BinaryOp { op: Arg::marker() }, - // STORE_SUBSCR specializations - Self::StoreSubscrDict | Self::StoreSubscrListInt => Self::StoreSubscr, - // SEND specializations - Self::SendGen => Self::Send { - delta: Arg::marker(), - }, - // UNPACK_SEQUENCE specializations + | Self::BinaryOpInplaceAddUnicode => Opcode::BinaryOp, + Self::StoreSubscrDict | Self::StoreSubscrListInt => Opcode::StoreSubscr, + Self::SendGen => Opcode::Send, Self::UnpackSequenceTwoTuple | Self::UnpackSequenceTuple | Self::UnpackSequenceList => { - Self::UnpackSequence { - count: Arg::marker(), - } + Opcode::UnpackSequence } - // STORE_ATTR specializations + Self::StoreAttrInstanceValue | Self::StoreAttrSlot | Self::StoreAttrWithHint => { - Self::StoreAttr { - namei: Arg::marker(), - } + Opcode::StoreAttr } - // LOAD_GLOBAL specializations - Self::LoadGlobalModule | Self::LoadGlobalBuiltin => Self::LoadGlobal { - namei: Arg::marker(), - }, - // LOAD_SUPER_ATTR specializations - Self::LoadSuperAttrAttr | Self::LoadSuperAttrMethod => Self::LoadSuperAttr { - namei: Arg::marker(), - }, - // LOAD_ATTR specializations + Self::LoadGlobalModule | Self::LoadGlobalBuiltin => Opcode::LoadGlobal, + Self::LoadSuperAttrAttr | Self::LoadSuperAttrMethod => Opcode::LoadSuperAttr, Self::LoadAttrInstanceValue | Self::LoadAttrModule | Self::LoadAttrWithHint @@ -685,28 +661,13 @@ impl Instruction { | Self::LoadAttrMethodNoDict | Self::LoadAttrMethodLazyDict | Self::LoadAttrNondescriptorWithValues - | Self::LoadAttrNondescriptorNoDict => Self::LoadAttr { - namei: Arg::marker(), - }, - // COMPARE_OP specializations - Self::CompareOpFloat | Self::CompareOpInt | Self::CompareOpStr => Self::CompareOp { - opname: Arg::marker(), - }, - // CONTAINS_OP specializations - Self::ContainsOpSet | Self::ContainsOpDict => Self::ContainsOp { - invert: Arg::marker(), - }, - // JUMP_BACKWARD specializations - Self::JumpBackwardNoJit | Self::JumpBackwardJit => Self::JumpBackward { - delta: Arg::marker(), - }, - // FOR_ITER specializations + | Self::LoadAttrNondescriptorNoDict => Opcode::LoadAttr, + Self::CompareOpFloat | Self::CompareOpInt | Self::CompareOpStr => Opcode::CompareOp, + Self::ContainsOpSet | Self::ContainsOpDict => Opcode::ContainsOp, + Self::JumpBackwardNoJit | Self::JumpBackwardJit => Opcode::JumpBackward, Self::ForIterList | Self::ForIterTuple | Self::ForIterRange | Self::ForIterGen => { - Self::ForIter { - delta: Arg::marker(), - } + Opcode::ForIter } - // CALL specializations Self::CallBoundMethodExactArgs | Self::CallPyExactArgs | Self::CallType1 @@ -726,15 +687,12 @@ impl Instruction { | Self::CallAllocAndEnterInit | Self::CallPyGeneral | Self::CallBoundMethodGeneral - | Self::CallNonPyGeneral => Self::Call { - argc: Arg::marker(), - }, - // CALL_KW specializations - Self::CallKwBoundMethod | Self::CallKwPy | Self::CallKwNonPy => Self::CallKw { - argc: Arg::marker(), - }, + | Self::CallNonPyGeneral => Opcode::Call, + Self::CallKwBoundMethod | Self::CallKwPy | Self::CallKwNonPy => Opcode::CallKw, _ => return None, - }) + }; + + Some(opcode.as_instruction()) } /// Map a specialized or instrumented opcode back to its adaptive (base) variant. @@ -1447,11 +1405,11 @@ impl InstructionMetadata for PseudoInstruction { /// SETUP_FINALLY: +1 (exc) /// SETUP_CLEANUP: +2 (lasti + exc) /// SETUP_WITH: +1 (pops __enter__ result, pushes lasti + exc) - fn stack_effect_jump(&self, _oparg: u32) -> i32 { + fn stack_effect_jump(&self, oparg: u32) -> i32 { match self { Self::SetupFinally { .. } | Self::SetupWith { .. } => 1, Self::SetupCleanup { .. } => 2, - _ => self.stack_effect(_oparg), + _ => self.stack_effect(oparg), } } diff --git a/crates/compiler-core/src/bytecode/oparg.rs b/crates/compiler-core/src/bytecode/oparg.rs index 0a2a660a4b3..7a9df84489e 100644 --- a/crates/compiler-core/src/bytecode/oparg.rs +++ b/crates/compiler-core/src/bytecode/oparg.rs @@ -19,6 +19,18 @@ impl OpArgByte { pub const fn new(value: u8) -> Self { Self(value) } + + /// Returns the inner value as a [`u8`]. + #[must_use] + pub const fn as_u8(self) -> u8 { + self.0 + } + + /// Returns the inner value as a [`u32`]. + #[must_use] + pub const fn as_u32(self) -> u32 { + self.0 as u32 + } } impl From for OpArgByte { @@ -29,7 +41,7 @@ impl From for OpArgByte { impl From for u8 { fn from(value: OpArgByte) -> Self { - value.0 + value.as_u8() } } @@ -93,7 +105,7 @@ pub struct OpArgState { impl OpArgState { #[inline(always)] - pub fn get(&mut self, ins: CodeUnit) -> (Instruction, OpArg) { + pub const fn get(&mut self, ins: CodeUnit) -> (Instruction, OpArg) { let arg = self.extend(ins.arg); if !matches!(ins.op, Instruction::ExtendedArg) { self.reset(); @@ -102,9 +114,9 @@ impl OpArgState { } #[inline(always)] - pub fn extend(&mut self, arg: OpArgByte) -> OpArg { - self.state = (self.state << 8) | u32::from(arg.0); - self.state.into() + pub const fn extend(&mut self, arg: OpArgByte) -> OpArg { + self.state = (self.state << 8) | arg.as_u32(); + OpArg::new(self.state) } #[inline(always)] @@ -124,10 +136,6 @@ impl OpArgState { /// - impl [`Into`] /// - impl [`OpArgType`] /// -/// # Note -/// If an enum variant has "alternative" values (i.e. `Foo = 0 | 1`), the first value will be the -/// result of converting to a number. -/// /// # Examples /// /// ```ignore @@ -138,8 +146,8 @@ impl OpArgState { /// /// Some doc. /// Foo = 4, /// Bar = 8, -/// Baz = 15 | 16, -/// Qux = 23 | 42 +/// Baz = 15, +/// Qux = 16 /// } /// ); /// ``` @@ -149,7 +157,7 @@ macro_rules! oparg_enum { $vis:vis enum $name:ident { $( $(#[$variant_meta:meta])* - $variant:ident = $value:literal $(| $alternatives:expr)* + $variant:ident = $value:literal ),* $(,)? } ) => { @@ -162,9 +170,9 @@ macro_rules! oparg_enum { } impl_oparg_enum!( - enum $name { + $vis enum $name { $( - $variant = $value $(| $alternatives)*, + $variant = $value, )* } ); @@ -173,48 +181,73 @@ macro_rules! oparg_enum { macro_rules! impl_oparg_enum { ( - enum $name:ident { + $vis:vis enum $name:ident { $( - $variant:ident = $value:literal $(| $alternatives:expr)* + $variant:ident = $value:literal ),* $(,)? } ) => { - impl TryFrom for $name { - type Error = $crate::marshal::MarshalError; + impl $name { + /// Returns the oparg as a [`u8`] value. + #[must_use] + $vis const fn as_u8(self) -> u8 { + match self { + $( + Self::$variant => $value, + )* + } + } - fn try_from(value: u8) -> Result { + /// Returns the oparg as a [`u32`] value. + #[must_use] + $vis const fn as_u32(self) -> u32 { + self.as_u8() as u32 + } + + $vis const fn try_from_u8(value: u8) -> Result { Ok(match value { $( - $value $(| $alternatives)* => Self::$variant, + $value => Self::$variant, )* - _ => return Err(Self::Error::InvalidBytecode), + _ => return Err($crate::marshal::MarshalError::InvalidBytecode), }) } + + $vis const fn try_from_u32(value: u32) -> Result { + if value > (u8::MAX as u32) { + return Err($crate::marshal::MarshalError::InvalidBytecode); + } + + // We already validated this is a lossles cast. + Self::try_from_u8(value as u8) + } + } + + impl TryFrom for $name { + type Error = $crate::marshal::MarshalError; + + fn try_from(value: u8) -> Result { + Self::try_from_u8(value) + } } impl TryFrom for $name { type Error = $crate::marshal::MarshalError; fn try_from(value: u32) -> Result { - u8::try_from(value) - .map_err(|_| Self::Error::InvalidBytecode) - .map(TryInto::try_into)? + Self::try_from_u32(value) } } impl From<$name> for u8 { fn from(value: $name) -> Self { - match value { - $( - $name::$variant => $value, - )* - } + value.as_u8() } } impl From<$name> for u32 { fn from(value: $name) -> Self { - Self::from(u8::from(value)) + value.as_u32() } } @@ -227,7 +260,7 @@ oparg_enum!( /// /// ## See also /// - /// - [CPython FVC_* flags](https://github.com/python/cpython/blob/8183fa5e3f78ca6ab862de7fb8b14f3d929421e0/Include/ceval.h#L129-L132) + /// - [CPython FVC_* flags](https://github.com/python/cpython/blob/v3.14.4/Include/ceval.h#L129-L132) #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] pub enum ConvertValueOparg { /// No conversion. @@ -236,8 +269,7 @@ oparg_enum!( /// f"{x}" /// f"{x:4}" /// ``` - // Ruff `ConversionFlag::None` is `-1i8`, when its converted to `u8` its value is `u8::MAX`. - None = 0 | 255, + None = 0, /// Converts by calling `str()`. /// /// ```python @@ -735,19 +767,19 @@ macro_rules! newtype_oparg { $vis struct $name(u32); impl $name { - /// Creates a new [`$name`] instance. + #[doc = concat!("Creates a new [`", stringify!($name), "`] instance.")] #[must_use] pub const fn from_u32(value: u32) -> Self { Self(value) } - /// Returns the oparg as a `u32` value. + /// Returns the oparg as a [`u32`] value. #[must_use] pub const fn as_u32(self) -> u32 { self.0 } - /// Returns the oparg as a `usize` value. + /// Returns the oparg as a [`usize`] value. #[must_use] pub const fn as_usize(self) -> usize { self.0 as usize diff --git a/crates/sre_engine/Cargo.toml b/crates/sre_engine/Cargo.toml index 4f899e6b3e9..8400a34b567 100644 --- a/crates/sre_engine/Cargo.toml +++ b/crates/sre_engine/Cargo.toml @@ -19,6 +19,7 @@ rustpython-wtf8 = { workspace = true } num_enum = { workspace = true } bitflags = { workspace = true } optional = { workspace = true } +icu_properties = { workspace = true } [dev-dependencies] criterion = { workspace = true } diff --git a/crates/sre_engine/src/string.rs b/crates/sre_engine/src/string.rs index 489819bfb3e..1350c9a07e1 100644 --- a/crates/sre_engine/src/string.rs +++ b/crates/sre_engine/src/string.rs @@ -1,3 +1,4 @@ +use icu_properties::props::{EnumeratedProperty, GeneralCategory, GeneralCategoryGroup}; use rustpython_wtf8::Wtf8; #[derive(Debug, Clone, Copy)] @@ -443,7 +444,11 @@ pub(crate) const fn is_uni_linebreak(ch: u32) -> bool { pub(crate) fn is_uni_alnum(ch: u32) -> bool { // TODO: check with cpython char::try_from(ch) - .map(|x| x.is_alphanumeric()) + .map(|c| { + GeneralCategoryGroup::Letter + .union(GeneralCategoryGroup::Number) + .contains(GeneralCategory::for_char(c)) + }) .unwrap_or(false) } #[inline] diff --git a/crates/stdlib/src/openssl.rs b/crates/stdlib/src/openssl.rs index 00461f1b468..88e616aa5d3 100644 --- a/crates/stdlib/src/openssl.rs +++ b/crates/stdlib/src/openssl.rs @@ -59,7 +59,8 @@ mod _ssl { }; use crate::{ common::lock::{ - PyMappedRwLockReadGuard, PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard, + LazyLock, PyMappedRwLockReadGuard, PyMutex, PyRwLock, PyRwLockReadGuard, + PyRwLockWriteGuard, }, socket::{self, PySocket}, vm::{ @@ -1046,7 +1047,7 @@ mod _ssl { #[pymethod] fn set_ciphers(&self, cipherlist: PyStrRef, vm: &VirtualMachine) -> PyResult<()> { - let ciphers = cipherlist.as_str(); + let ciphers: &str = cipherlist.as_ref(); if ciphers.contains('\0') { return Err(exceptions::cstring_error(vm)); } @@ -1102,7 +1103,8 @@ mod _ssl { // Convert name to CString, supporting both str and bytes let name_cstr = match name { Either::A(s) => { - if s.as_str().contains('\0') { + let s: &str = s.as_ref(); + if s.contains('\0') { return Err(exceptions::cstring_error(vm)); } s.to_cstring(vm)? @@ -1459,7 +1461,7 @@ mod _ssl { } *self.psk_server_callback.lock() = Some(callback); if let OptionalArg::Present(hint) = identity_hint { - *self.psk_identity_hint.lock() = Some(hint.as_str().to_owned()); + *self.psk_identity_hint.lock() = Some(hint.to_string()); } // Note: The actual callback will be invoked via SSL app_data mechanism } @@ -1487,7 +1489,8 @@ mod _ssl { if let Some(cadata) = args.cadata { let (certs, is_pem) = match cadata { Either::A(s) => { - if !s.as_str().is_ascii() { + let s: &str = s.as_ref(); + if !s.is_ascii() { return Err(invalid_cadata(vm)); } (X509::stack_from_pem(s.as_bytes()), true) @@ -1977,7 +1980,7 @@ mod _ssl { use crate::vm::builtins::{PyByteArray, PyBytes, PyStr}; if let Some(s) = obj.downcast_ref::() { - Ok(s.as_str().as_bytes().to_vec()) + Ok(s.as_bytes().to_vec()) } else if let Some(b) = obj.downcast_ref::() { Ok(b.as_bytes().to_vec()) } else if let Some(ba) = obj.downcast_ref::() { @@ -2033,7 +2036,7 @@ mod _ssl { // Configure server hostname if let Some(hostname) = &server_hostname { - let hostname_str = hostname.as_str(); + let hostname_str: &str = hostname.as_ref(); if hostname_str.is_empty() || hostname_str.starts_with('.') { return Err(vm.new_value_error( "server_hostname cannot be an empty string or start with a leading dot.", @@ -2752,7 +2755,7 @@ mod _ssl { ) -> PyResult> { const CB_MAXLEN: usize = 512; - let cb_type_str = cb_type.as_ref().map_or("tls-unique", |s| s.as_str()); + let cb_type_str = cb_type.as_ref().map_or("tls-unique", |s| s.as_ref()); if cb_type_str != "tls-unique" { return Err(vm.new_value_error(format!( diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 0ddbd9b8513..aafbe196a07 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -1,6 +1,6 @@ //! Infamous code object. The python class `code` -use super::{PyBytesRef, PyStrRef, PyTupleRef, PyType}; +use super::{PyBytesRef, PyStrRef, PyTupleRef, PyType, set::PyFrozenSet}; use crate::common::lock::PyMutex; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, @@ -267,13 +267,6 @@ impl<'a> AsBag for &'a Context { } } -impl<'a> AsBag for &'a VirtualMachine { - type Bag = PyObjBag<'a>; - fn as_bag(self) -> PyObjBag<'a> { - PyObjBag(&self.ctx) - } -} - #[derive(Clone, Copy)] pub struct PyObjBag<'a>(pub &'a Context); @@ -348,6 +341,87 @@ impl ConstantBag for PyObjBag<'_> { } } +#[derive(Clone, Copy)] +pub struct PyVmBag<'a>(pub &'a VirtualMachine); + +impl ConstantBag for PyVmBag<'_> { + type Constant = Literal; + + fn make_constant(&self, constant: BorrowedConstant<'_, C>) -> Self::Constant { + let vm = self.0; + let ctx = &vm.ctx; + let obj = match constant { + BorrowedConstant::Integer { value } => ctx.new_bigint(value).into(), + BorrowedConstant::Float { value } => ctx.new_float(value).into(), + BorrowedConstant::Complex { value } => ctx.new_complex(value).into(), + BorrowedConstant::Str { value } if value.len() <= 20 => { + ctx.intern_str(value).to_object() + } + BorrowedConstant::Str { value } => ctx.new_str(value).into(), + BorrowedConstant::Bytes { value } => ctx.new_bytes(value.to_vec()).into(), + BorrowedConstant::Boolean { value } => ctx.new_bool(value).into(), + BorrowedConstant::Code { code } => { + PyCode::new_ref_with_bag(vm, code.map_clone_bag(self)).into() + } + BorrowedConstant::Tuple { elements } => { + let elements = elements + .iter() + .map(|constant| self.make_constant(constant.borrow_constant()).0) + .collect(); + ctx.new_tuple(elements).into() + } + BorrowedConstant::Slice { elements } => { + let [start, stop, step] = elements; + let start_obj = self.make_constant(start.borrow_constant()).0; + let stop_obj = self.make_constant(stop.borrow_constant()).0; + let step_obj = self.make_constant(step.borrow_constant()).0; + use crate::builtins::PySlice; + PySlice { + start: Some(start_obj), + stop: stop_obj, + step: Some(step_obj), + } + .into_ref(ctx) + .into() + } + BorrowedConstant::Frozenset { elements } => { + let elements = elements + .iter() + .map(|constant| self.make_constant(constant.borrow_constant()).0); + PyFrozenSet::from_iter(vm, elements) + .unwrap() + .into_ref(ctx) + .into() + } + BorrowedConstant::None => ctx.none(), + BorrowedConstant::Ellipsis => ctx.ellipsis.clone().into(), + }; + + Literal(obj) + } + + fn make_name(&self, name: &str) -> &'static PyStrInterned { + self.0.ctx.intern_str(name) + } + + fn make_int(&self, value: BigInt) -> Self::Constant { + Literal(self.0.ctx.new_int(value).into()) + } + + fn make_tuple(&self, elements: impl Iterator) -> Self::Constant { + Literal( + self.0 + .ctx + .new_tuple(elements.map(|lit| lit.0).collect()) + .into(), + ) + } + + fn make_code(&self, code: CodeObject) -> Self::Constant { + Literal(PyCode::new_ref_with_bag(self.0, code).into()) + } +} + pub type CodeObject = bytecode::CodeObject; pub trait IntoCodeObject { @@ -427,6 +501,22 @@ impl PyCode { Ordering::Relaxed, ); } + + pub fn new_ref_with_bag(vm: &VirtualMachine, code: CodeObject) -> PyRef { + PyRef::new_ref(PyCode::new(code), vm.ctx.types.code_type.to_owned(), None) + } + + pub fn new_ref_from_bytecode(vm: &VirtualMachine, code: bytecode::CodeObject) -> PyRef { + Self::new_ref_with_bag(vm, code.map_bag(PyVmBag(vm))) + } + + pub fn new_ref_from_frozen>( + vm: &VirtualMachine, + code: frozen::FrozenCodeObject, + ) -> PyRef { + Self::new_ref_with_bag(vm, code.decode(PyVmBag(vm))) + } + pub fn from_pyc_path(path: &std::path::Path, vm: &VirtualMachine) -> PyResult> { let name = match path.file_stem() { Some(stem) => stem.display().to_string(), @@ -1379,7 +1469,7 @@ impl ToPyObject for CodeObject { impl ToPyObject for bytecode::CodeObject { fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { - vm.ctx.new_code(self).into() + PyCode::new_ref_from_bytecode(vm, self).into() } } diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs index eae4ec7fd6b..9791435a23a 100644 --- a/crates/vm/src/builtins/dict.rs +++ b/crates/vm/src/builtins/dict.rs @@ -663,7 +663,8 @@ impl Py { ) -> PyResult> { if self.exact_dict(vm) { self.entries.get(vm, key) - // FIXME: check __missing__? + // Match CPython's exact-dict fast path: __missing__ only participates + // for dict subclasses through the generic mapping lookup path below. } else { match self.as_object().get_item(key, vm) { Ok(value) => Ok(Some(value)), diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index 870d3b72a74..af1c4a5ae92 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -45,7 +45,8 @@ use rustpython_common::{ }; use icu_properties::props::{ - BidiClass, BinaryProperty, EnumeratedProperty, GeneralCategory, XidContinue, XidStart, + BidiClass, BinaryProperty, EnumeratedProperty, GeneralCategory, GeneralCategoryGroup, + NumericType, XidContinue, XidStart, }; use unicode_casing::CharExt; @@ -946,21 +947,32 @@ impl PyStr { #[pymethod] fn isalnum(&self) -> bool { - !self.data.is_empty() && self.char_all(char::is_alphanumeric) + !self.data.is_empty() + && self.char_all(|c| { + GeneralCategoryGroup::Letter + .union(GeneralCategoryGroup::Number) + .contains(GeneralCategory::for_char(c)) + }) } #[pymethod] fn isnumeric(&self) -> bool { - !self.data.is_empty() && self.char_all(char::is_numeric) + !self.data.is_empty() + && self.char_all(|c| { + [ + NumericType::Decimal, + NumericType::Digit, + NumericType::Numeric, + ] + .contains(&NumericType::for_char(c)) + }) } #[pymethod] fn isdigit(&self) -> bool { - // python's isdigit also checks if exponents are digits, these are the unicode codepoints for exponents !self.data.is_empty() && self.char_all(|c| { - c.is_ascii_digit() - || matches!(c, '⁰' | '¹' | '²' | '³' | '⁴' | '⁵' | '⁶' | '⁷' | '⁸' | '⁹') + [NumericType::Digit, NumericType::Decimal].contains(&NumericType::for_char(c)) }) } @@ -1059,7 +1071,9 @@ impl PyStr { #[pymethod] fn isalpha(&self) -> bool { - !self.data.is_empty() && self.char_all(char::is_alphabetic) + !self.data.is_empty() + && self + .char_all(|c| GeneralCategoryGroup::Letter.contains(GeneralCategory::for_char(c))) } #[pymethod] diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 92e0a11d0f5..180c4fad0ed 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -2725,16 +2725,7 @@ impl ExecutingFrame<'_> { let class_dict = self.pop_value(); let idx = i.get(arg).as_usize(); let name = self.localsplus_name(idx); - // Only treat KeyError as "not found", propagate other exceptions - let value = if let Some(dict_obj) = class_dict.downcast_ref::() { - dict_obj.get_item_opt(name, vm)? - } else { - match class_dict.get_item(name, vm) { - Ok(v) => Some(v), - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => None, - Err(e) => return Err(e), - } - }; + let value = self.mapping_get_optional(&class_dict, name, vm)?; self.push_value(match value { Some(v) => v, None => self @@ -2748,18 +2739,7 @@ impl ExecutingFrame<'_> { // PEP 649: Pop dict from stack (classdict), check there first, then globals let dict = self.pop_value(); let name = self.code.names[idx.get(arg) as usize]; - - // Only treat KeyError as "not found", propagate other exceptions - let value = if let Some(dict_obj) = dict.downcast_ref::() { - dict_obj.get_item_opt(name, vm)? - } else { - // Not an exact dict, use mapping protocol - match dict.get_item(name, vm) { - Ok(v) => Some(v), - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => None, - Err(e) => return Err(e), - } - }; + let value = self.mapping_get_optional(&dict, name, vm)?; self.push_value(match value { Some(v) => v, @@ -3451,15 +3431,16 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::StoreFastLoadFast { var_nums } => { - let value = self.pop_value(); - let locals = self.localsplus.fastlocals_mut(); + // pop_value_opt: allows NULL from LoadFastAndClear restore paths. + let value = self.pop_value_opt(); let oparg = var_nums.get(arg); let (store_idx, load_idx) = oparg.indexes(); - locals[store_idx] = Some(value); - let load_value = locals[load_idx] - .clone() - .expect("StoreFastLoadFast: load slot should have value after store"); - self.push_value(load_value); + let load_value = { + let locals = self.localsplus.fastlocals_mut(); + locals[store_idx] = value; + locals[load_idx].clone() + }; + self.push_value_opt(load_value); Ok(None) } Instruction::StoreFastStoreFast { var_nums } => { @@ -6117,6 +6098,27 @@ impl ExecutingFrame<'_> { } } + #[inline] + fn mapping_get_optional( + &self, + mapping: &PyObjectRef, + name: &Py, + vm: &VirtualMachine, + ) -> PyResult> { + if mapping.class().is(vm.ctx.types.dict_type) { + let dict = mapping + .downcast_ref::() + .expect("exact dict must have a PyDict payload"); + dict.get_item_opt(name, vm) + } else { + match mapping.get_item(name, vm) { + Ok(value) => Ok(Some(value)), + Err(err) if err.fast_isinstance(vm.ctx.exceptions.key_error) => Ok(None), + Err(err) => Err(err), + } + } + } + #[inline] fn load_global_or_builtin(&self, name: &Py, vm: &VirtualMachine) -> PyResult { if let Some(builtins_dict) = self.builtins_dict { diff --git a/crates/vm/src/import.rs b/crates/vm/src/import.rs index 9d015c8f3b6..f8f41c12081 100644 --- a/crates/vm/src/import.rs +++ b/crates/vm/src/import.rs @@ -72,7 +72,7 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult> { vm.ctx.new_utf8_str(name), ) })?; - Ok(vm.ctx.new_code(frozen.code)) + Ok(PyCode::new_ref_from_frozen(vm, frozen.code)) } pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult { @@ -82,7 +82,12 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult { vm.ctx.new_utf8_str(module_name), ) })?; - let module = import_code_obj(vm, module_name, vm.ctx.new_code(frozen.code), false)?; + let module = import_code_obj( + vm, + module_name, + PyCode::new_ref_from_frozen(vm, frozen.code), + false, + )?; debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok()); let origname = resolve_frozen_alias(module_name); module.set_attr("__origname__", vm.ctx.new_utf8_str(origname), vm)?; diff --git a/crates/vm/src/stdlib/_ast.rs b/crates/vm/src/stdlib/_ast.rs index 73819e257c1..bde6916a663 100644 --- a/crates/vm/src/stdlib/_ast.rs +++ b/crates/vm/src/stdlib/_ast.rs @@ -776,7 +776,7 @@ pub(crate) fn compile( let source_file = SourceFileBuilder::new(filename, text).finish(); let code = codegen::compile::compile_top(ast, source_file, mode, opts) .map_err(|err| vm.new_syntax_error(&err.into(), None))?; // FIXME source - Ok(vm.ctx.new_code(code).into()) + Ok(crate::builtins::PyCode::new_ref_from_bytecode(vm, code).into()) } #[cfg(feature = "codegen")] diff --git a/crates/vm/src/stdlib/_ast/node.rs b/crates/vm/src/stdlib/_ast/node.rs index 6f74e71511f..4ee3893b665 100644 --- a/crates/vm/src/stdlib/_ast/node.rs +++ b/crates/vm/src/stdlib/_ast/node.rs @@ -31,7 +31,14 @@ impl Node for Vec { source_file: &SourceFile, object: PyObjectRef, ) -> PyResult { - vm.extract_elements_with(&object, |obj| Node::ast_from_object(vm, source_file, obj)) + // Recursion guard for each element: prevents stack overflow when a + // sequence element transitively references the sequence itself + // (e.g. `l = ast.List(...); l.elts = [l]`). See issue #4862. + vm.extract_elements_with(&object, |obj| { + vm.with_recursion("while traversing AST node", || { + Node::ast_from_object(vm, source_file, obj) + }) + }) } } @@ -45,7 +52,13 @@ impl Node for Box { source_file: &SourceFile, object: PyObjectRef, ) -> PyResult { - T::ast_from_object(vm, source_file, object).map(Self::new) + // Recursion guard: every descent through a Box increments the + // VM's recursion depth so cyclic or pathologically deep ASTs raise + // RecursionError instead of overflowing the native stack. + // See issue #4862. + vm.with_recursion("while traversing AST node", || { + T::ast_from_object(vm, source_file, object).map(Self::new) + }) } fn is_none(&self) -> bool { diff --git a/crates/vm/src/stdlib/_ctypes/simple.rs b/crates/vm/src/stdlib/_ctypes/simple.rs index 8f99fa8e57a..3bf1f84fbc5 100644 --- a/crates/vm/src/stdlib/_ctypes/simple.rs +++ b/crates/vm/src/stdlib/_ctypes/simple.rs @@ -309,7 +309,7 @@ impl PyCSimpleType { // Float types: accept numbers Some(tc @ ("f" | "d" | "g")) - if (value.try_float(vm).is_ok() || value.try_int(vm).is_ok()) => + if value.try_float(vm).is_ok() || value.try_int(vm).is_ok() => { return create_simple_with_value(tc, &value); } diff --git a/crates/vm/src/stdlib/_imp.rs b/crates/vm/src/stdlib/_imp.rs index c0acb304a64..71a8c091e5e 100644 --- a/crates/vm/src/stdlib/_imp.rs +++ b/crates/vm/src/stdlib/_imp.rs @@ -261,11 +261,11 @@ mod _imp { name.clone().into_wtf8(), ) }; - let bag = crate::builtins::code::PyObjBag(&vm.ctx); + let bag = crate::builtins::code::PyVmBag(vm); let code = rustpython_compiler_core::marshal::deserialize_code(&mut &contiguous[..], bag) .map_err(|_| invalid_err())?; - return Ok(vm.ctx.new_code(code)); + return Ok(PyCode::new_ref_with_bag(vm, code)); } import::make_frozen(vm, name.as_str()) } diff --git a/crates/vm/src/stdlib/marshal.rs b/crates/vm/src/stdlib/marshal.rs index dace6bbf3e3..b19cc5eb52e 100644 --- a/crates/vm/src/stdlib/marshal.rs +++ b/crates/vm/src/stdlib/marshal.rs @@ -3,7 +3,7 @@ pub(crate) use decl::module_def; #[pymodule(name = "marshal")] mod decl { - use crate::builtins::code::{CodeObject, Literal, PyObjBag}; + use crate::builtins::code::{CodeObject, Literal, PyVmBag}; use crate::class::StaticType; use crate::common::wtf8::Wtf8; use crate::{ @@ -382,7 +382,7 @@ mod decl { impl<'a> marshal::MarshalBag for PyMarshalBag<'a> { type Value = PyObjectRef; - type ConstantBag = PyObjBag<'a>; + type ConstantBag = PyVmBag<'a>; fn make_bool(&self, value: bool) -> Self::Value { self.0.ctx.new_bool(value).into() @@ -412,7 +412,7 @@ mod decl { self.0.ctx.new_tuple(elements.collect()).into() } fn make_code(&self, code: CodeObject) -> Self::Value { - self.0.ctx.new_code(code).into() + crate::builtins::PyCode::new_ref_with_bag(self.0, code).into() } fn make_stop_iter(&self) -> Result { Ok(self.0.ctx.exceptions.stop_iteration.to_owned().into()) @@ -472,7 +472,7 @@ mod decl { .into()) } fn constant_bag(self) -> Self::ConstantBag { - PyObjBag(&self.0.ctx) + PyVmBag(self.0) } } diff --git a/crates/vm/src/stdlib/time.rs b/crates/vm/src/stdlib/time.rs index 2cc2f5f50b8..3f2c3ba8288 100644 --- a/crates/vm/src/stdlib/time.rs +++ b/crates/vm/src/stdlib/time.rs @@ -27,7 +27,7 @@ unsafe extern "C" { mod decl { use crate::{ AsObject, Py, PyObjectRef, PyResult, VirtualMachine, - builtins::{PyStrRef, PyTypeRef}, + builtins::{PyBaseExceptionRef, PyStrRef, PyTypeRef}, function::{Either, FuncArgs, OptionalArg}, types::{PyStructSequence, struct_sequence_new}, }; @@ -354,6 +354,13 @@ mod decl { ) -> PyResult { let invalid_tuple = || vm.new_type_error(format!("{func_name}(): illegal time tuple argument")); + let classify_err = |e: PyBaseExceptionRef| { + if e.class().is(vm.ctx.exceptions.overflow_error) { + vm.new_overflow_error(format!("{func_name} argument out of range")) + } else { + invalid_tuple() + } + }; let year: i64 = t.tm_year.clone().try_into_value(vm).map_err(|e| { if e.class().is(vm.ctx.exceptions.overflow_error) { @@ -370,46 +377,30 @@ mod decl { .tm_mon .clone() .try_into_value::(vm) - .map_err(|_| invalid_tuple())? + .map_err(classify_err)? - 1; - let tm_mday = t - .tm_mday - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - let tm_hour = t - .tm_hour - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - let tm_min = t - .tm_min - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - let tm_sec = t - .tm_sec - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + let tm_mday = t.tm_mday.clone().try_into_value(vm).map_err(classify_err)?; + let tm_hour = t.tm_hour.clone().try_into_value(vm).map_err(classify_err)?; + let tm_min = t.tm_min.clone().try_into_value(vm).map_err(classify_err)?; + let tm_sec = t.tm_sec.clone().try_into_value(vm).map_err(classify_err)?; let tm_wday = (t .tm_wday .clone() .try_into_value::(vm) - .map_err(|_| invalid_tuple())? + .map_err(classify_err)? + 1) % 7; let tm_yday = t .tm_yday .clone() .try_into_value::(vm) - .map_err(|_| invalid_tuple())? + .map_err(classify_err)? - 1; let tm_isdst = t .tm_isdst .clone() .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + .map_err(classify_err)?; let mut tm: libc::tm = unsafe { core::mem::zeroed() }; tm.tm_year = year - 1900; @@ -474,7 +465,7 @@ mod decl { .tm_gmtoff .clone() .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + .map_err(classify_err)?; tm.tm_gmtoff = gmtoff as _; } @@ -963,41 +954,28 @@ mod decl { vm: &VirtualMachine, ) -> PyResult { let invalid_tuple = || vm.new_type_error("mktime(): illegal time tuple argument"); - let year: i32 = t - .tm_year - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + let classify_err = |e: PyBaseExceptionRef| { + if e.class().is(vm.ctx.exceptions.overflow_error) { + vm.new_overflow_error("mktime argument out of range") + } else { + invalid_tuple() + } + }; + let year: i32 = t.tm_year.clone().try_into_value(vm).map_err(classify_err)?; if year < i32::MIN + 1900 { return Err(vm.new_overflow_error("year out of range")); } let mut tm: libc::tm = unsafe { core::mem::zeroed() }; - tm.tm_sec = t - .tm_sec - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - tm.tm_min = t - .tm_min - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - tm.tm_hour = t - .tm_hour - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; - tm.tm_mday = t - .tm_mday - .clone() - .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + tm.tm_sec = t.tm_sec.clone().try_into_value(vm).map_err(classify_err)?; + tm.tm_min = t.tm_min.clone().try_into_value(vm).map_err(classify_err)?; + tm.tm_hour = t.tm_hour.clone().try_into_value(vm).map_err(classify_err)?; + tm.tm_mday = t.tm_mday.clone().try_into_value(vm).map_err(classify_err)?; tm.tm_mon = t .tm_mon .clone() .try_into_value::(vm) - .map_err(|_| invalid_tuple())? + .map_err(classify_err)? - 1; tm.tm_year = year - 1900; tm.tm_wday = -1; @@ -1005,13 +983,13 @@ mod decl { .tm_yday .clone() .try_into_value::(vm) - .map_err(|_| invalid_tuple())? + .map_err(classify_err)? - 1; tm.tm_isdst = t .tm_isdst .clone() .try_into_value(vm) - .map_err(|_| invalid_tuple())?; + .map_err(classify_err)?; Ok(tm) } diff --git a/crates/vm/src/types/slot.rs b/crates/vm/src/types/slot.rs index 232b55110a2..31a03094e8f 100644 --- a/crates/vm/src/types/slot.rs +++ b/crates/vm/src/types/slot.rs @@ -1394,9 +1394,9 @@ impl PyType { SlotAccessor::SqLength => { update_sub_slot!(as_sequence, length, sequence_len_wrapper, SeqLength) } - // Sequence concat uses sq_concat slot - no generic wrapper needed - // (handled by number protocol fallback) SlotAccessor::SqConcat | SlotAccessor::SqInplaceConcat if !ADD => { + // Sequence concat uses sq_concat slot - no generic wrapper needed + // (handled by number protocol fallback) accessor.inherit_from_mro(self); } SlotAccessor::SqRepeat => { diff --git a/crates/vm/src/vm/compile.rs b/crates/vm/src/vm/compile.rs index 7294dc8f897..221df849f62 100644 --- a/crates/vm/src/vm/compile.rs +++ b/crates/vm/src/vm/compile.rs @@ -25,8 +25,8 @@ impl VirtualMachine { source_path: String, opts: CompileOpts, ) -> Result, CompileError> { - let code = - compiler::compile(source, mode, &source_path, opts).map(|code| self.ctx.new_code(code)); + let code = compiler::compile(source, mode, &source_path, opts) + .map(|code| PyCode::new_ref_from_bytecode(self, code)); #[cfg(feature = "parser")] if code.is_ok() { self.emit_string_escape_warnings(source, &source_path); diff --git a/crates/vm/src/vm/vm_ops.rs b/crates/vm/src/vm/vm_ops.rs index abce0ea678b..9cc2853a09a 100644 --- a/crates/vm/src/vm/vm_ops.rs +++ b/crates/vm/src/vm/vm_ops.rs @@ -372,7 +372,7 @@ impl VirtualMachine { } else { self.new_type_error(format!( "unsupported operand type(s) for {}: \ - '{}' and '{}', '{}'", + '{}', '{}', '{}'", op_str, a.class(), b.class(), @@ -401,7 +401,7 @@ impl VirtualMachine { binary_func!(_sub, Subtract, "-"); binary_func!(_mod, Remainder, "%"); - binary_func!(_divmod, Divmod, "divmod"); + binary_func!(_divmod, Divmod, "divmod()"); binary_func!(_lshift, Lshift, "<<"); binary_func!(_rshift, Rshift, ">>"); binary_func!(_and, And, "&"); diff --git a/extra_tests/snippets/builtin_str.py b/extra_tests/snippets/builtin_str.py index 3d54643b3ce..3899c04956e 100644 --- a/extra_tests/snippets/builtin_str.py +++ b/extra_tests/snippets/builtin_str.py @@ -72,6 +72,28 @@ assert "\u1c89".istitle() # assert "DZ".title() == "Dz" assert a.isalpha() +assert not "\u093f".isalpha() + +# Combining characters differ slightly between Rust and Python +assert "\u006e".isalnum() +assert not "\u0303".isalnum() +assert not "\u006e\u0303".isalnum() +assert "\u00f1".isalnum() +assert not "\u0345".isalnum() +assert not "\u093f".isalnum() +for raw in range(0x0363, 0x036F): + assert not chr(raw).isalnum() + +# isdigit is true for exponents +assert "⁰".isdigit() +assert "⁰".isnumeric() +assert not "½".isdigit() +assert "½".isnumeric() +assert not "Ⅻ".isdigit() +assert "Ⅻ".isnumeric() + +# isnumeric is broader than Rust's +assert "\u3405".isnumeric() s = "1 2 3" assert s.split(" ", 1) == ["1", "2 3"] diff --git a/extra_tests/snippets/stdlib_ast.py b/extra_tests/snippets/stdlib_ast.py index dc626506fa9..7b5c69df49e 100644 --- a/extra_tests/snippets/stdlib_ast.py +++ b/extra_tests/snippets/stdlib_ast.py @@ -37,3 +37,54 @@ def foo(): assert i.module is None assert i.names[0].name == "a" assert i.names[0].asname is None + + +# Regression test for issue #4862: +# A cyclic AST fed to compile() used to overflow the Rust stack and SIGSEGV. +# After the fix, the recursion guard in ast_from_object raises RecursionError, +# matching CPython's behavior. Covers both Box descents (UnaryOp, BinOp, +# Call, Attribute) and Vec descents (List, Tuple). +import warnings + + +def _cyclic_cases(): + # Box descents + u = ast.UnaryOp(op=ast.Not(), lineno=0, col_offset=0) + u.operand = u + yield "UnaryOp", u + + b = ast.BinOp( + op=ast.Add(), + right=ast.Constant(value=0, lineno=0, col_offset=0), + lineno=0, + col_offset=0, + ) + b.left = b + yield "BinOp", b + + c = ast.Call(args=[], keywords=[], lineno=0, col_offset=0) + c.func = c + yield "Call", c + + a = ast.Attribute(attr="x", ctx=ast.Load(), lineno=0, col_offset=0) + a.value = a + yield "Attribute", a + + # Vec descents + lst = ast.List(ctx=ast.Load(), lineno=0, col_offset=0) + lst.elts = [lst] + yield "List", lst + + tup = ast.Tuple(ctx=ast.Load(), lineno=0, col_offset=0) + tup.elts = [tup] + yield "Tuple", tup + + +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + for name, node in _cyclic_cases(): + try: + compile(ast.Expression(node), "", "eval") + raise AssertionError(f"cyclic {name} should raise RecursionError") + except RecursionError: + pass # expected; matches CPython diff --git a/extra_tests/snippets/stdlib_re.py b/extra_tests/snippets/stdlib_re.py index 53f21f91734..8613ddd30fc 100644 --- a/extra_tests/snippets/stdlib_re.py +++ b/extra_tests/snippets/stdlib_re.py @@ -79,3 +79,6 @@ # Test of fix re.fullmatch POSSESSIVE_REPEAT, issue #7183 assert re.fullmatch(r"([0-9]++(?:\.[0-9]+)*+)", "1.25.38") assert re.fullmatch(r"([0-9]++(?:\.[0-9]+)*+)", "1.25.38").group(0) == "1.25.38" + +# Combining characters; issue #7518 +assert not re.match(r"\w", "\u0345"), r"\w should not match U+0345 (category Mn)" diff --git a/extra_tests/snippets/stdlib_time.py b/extra_tests/snippets/stdlib_time.py index 9a92969f5f9..1629443a9e7 100644 --- a/extra_tests/snippets/stdlib_time.py +++ b/extra_tests/snippets/stdlib_time.py @@ -17,3 +17,35 @@ s = time.asctime(x) # print(s) assert s == "Thu Jan 1 00:16:40 1970" + + +# Regression test for RustPython issue #4938: +# struct_time field overflow should raise OverflowError (matching CPython), +# not TypeError. Covers mktime, asctime, and strftime. +I32_MAX_PLUS_1 = 2147483648 +overflow_cases = [ + (I32_MAX_PLUS_1, 1, 1, 0, 0, 0, 0, 0, 0), # i32 overflow in year + (2024, I32_MAX_PLUS_1, 1, 0, 0, 0, 0, 0, 0), # i32 overflow in month + (2024, 1, I32_MAX_PLUS_1, 0, 0, 0, 0, 0, 0), # i32 overflow in mday + (2024, 1, 1, 0, 0, I32_MAX_PLUS_1, 0, 0, 0), # i32 overflow in sec + (88888888888,) * 9, # multi-field i32 overflow +] + +for case in overflow_cases: + for func_name, call in [ + ("mktime", lambda c=case: time.mktime(c)), + ("asctime", lambda c=case: time.asctime(c)), + ("strftime", lambda c=case: time.strftime("%Y", c)), + ]: + try: + call() + except OverflowError: + pass # expected, matches CPython + except TypeError as e: + raise AssertionError( + f"{func_name}({case}) raised TypeError (should be OverflowError): {e}" + ) from e + else: + raise AssertionError( + f"{func_name}({case}) did not raise — expected OverflowError" + ) diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index e30a8955fdd..e8b9c1bf5f8 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -11,6 +11,8 @@ """ import argparse +import ast +import builtins import dis import json import os @@ -38,6 +40,7 @@ "JUMP_IF_TRUE_OR_POP", "JUMP_IF_FALSE_OR_POP", "FOR_ITER", + "END_ASYNC_FOR", "SEND", } ) @@ -93,6 +96,14 @@ def _unescape(m): argrepr = re.sub(r"\\u([0-9a-fA-F]{4})", _unescape, argrepr) argrepr = re.sub(r"\\U([0-9a-fA-F]{8})", _unescape, argrepr) + if argrepr.startswith("frozenset({") and argrepr.endswith("})"): + try: + values = ast.literal_eval(argrepr[len("frozenset(") : -1]) + except Exception: + return argrepr + if isinstance(values, set): + parts = sorted(_normalize_argrepr(repr(value)) for value in values) + return f"frozenset({{{', '.join(parts)}}})" return argrepr @@ -100,10 +111,33 @@ def _unescape(m): hasattr(sys, "implementation") and sys.implementation.name == "rustpython" ) +if _IS_RUSTPYTHON and hasattr(dis, "_common_constants"): + common_constants = list(dis._common_constants) + while len(common_constants) < 7: + common_constants.append( + (builtins.list, builtins.set)[len(common_constants) - 5] + ) + dis._common_constants = common_constants + # RustPython's ComparisonOperator enum values → operator strings _RP_CMP_OPS = {0: "<", 1: "<=", 2: "==", 3: "!=", 4: ">", 5: ">="} +def _resolve_localsplus_name(code, arg): + if not isinstance(arg, int) or arg < 0: + return arg + nlocals = len(code.co_varnames) + if arg < nlocals: + return code.co_varnames[arg] + varnames_set = set(code.co_varnames) + nonparam_cells = [v for v in code.co_cellvars if v not in varnames_set] + extra = nonparam_cells + list(code.co_freevars) + idx = arg - nlocals + if 0 <= idx < len(extra): + return extra[idx] + return arg + + def _resolve_arg_fallback(code, opname, arg): """Resolve a raw argument to its human-readable form. @@ -113,8 +147,7 @@ def _resolve_arg_fallback(code, opname, arg): return arg try: if "FAST" in opname: - if 0 <= arg < len(code.co_varnames): - return code.co_varnames[arg] + return _resolve_localsplus_name(code, arg) elif opname == "LOAD_CONST": if 0 <= arg < len(code.co_consts): return _normalize_argrepr(repr(code.co_consts[arg])) @@ -125,18 +158,7 @@ def _resolve_arg_fallback(code, opname, arg): "LOAD_CLOSURE", "MAKE_CELL", ): - # arg is localsplus index: - # 0..nlocals-1 = varnames (parameter cells reuse these slots) - # nlocals.. = non-parameter cells + freevars - nlocals = len(code.co_varnames) - if arg < nlocals: - return code.co_varnames[arg] - varnames_set = set(code.co_varnames) - nonparam_cells = [v for v in code.co_cellvars if v not in varnames_set] - extra = nonparam_cells + list(code.co_freevars) - idx = arg - nlocals - if 0 <= idx < len(extra): - return extra[idx] + return _resolve_localsplus_name(code, arg) elif opname in ( "LOAD_NAME", "STORE_NAME", @@ -237,7 +259,7 @@ def _metadata_cache_slot_offsets(inst): # 1. argval not in offset_to_idx (not a valid byte offset) # 2. argval == arg (raw arg returned as-is, not resolved to offset) # 3. For backward jumps: argval should be < current offset - is_backward = "BACKWARD" in inst.opname + is_backward = "BACKWARD" in inst.opname or inst.opname == "END_ASYNC_FOR" argval_is_raw = inst.argval == inst.arg and inst.arg is not None if target_idx is None or argval_is_raw: target_idx = None # force recalculation diff --git a/scripts/generate_opcode_metadata.py b/scripts/generate_opcode_metadata.py index f16759ec87a..2758f0b982a 100644 --- a/scripts/generate_opcode_metadata.py +++ b/scripts/generate_opcode_metadata.py @@ -85,48 +85,44 @@ def extract_enum_body(text: str, name: str) -> str: return text[start + 1 : i] -def build_deopts(contents: str) -> dict[str, list[str]]: - raw_body = re.search( - r"fn deopt\(self\) -> Option(.*)", contents, re.DOTALL - ).group(1) - body = "\n".join( - itertools.takewhile( - lambda l: not l.startswith("_ =>"), # Take until reaching fallback - filter( - lambda l: ( - not l.startswith( - ("//", "Some(match") - ) # Skip comments or start of match - ), - map(str.strip, raw_body.splitlines()), - ), - ) - ).removeprefix("{") +def build_deopts(text: str) -> dict[str, list[str]]: + raw_body = re.search(r"fn deopt\(self\)(.*)", text, re.DOTALL).group(1) + match_start = raw_body.find("match self") + if match_start == -1: + raise ValueError("Could not detect a match statement in deopt method") - depth = 0 - arms = [] - buf = [] - for char in body: - if char == "{": - depth += 1 - elif char == "}": - depth -= 1 + brace_depth = 0 + block_start = None + block_end = None - if depth == 0 and (char in ("}", ",")): - arm = "".join(buf).strip() - arms.append(arm) - buf = [] - else: - buf.append(char) + for i, ch in enumerate(raw_body[match_start:], match_start): + if ch == "{": + brace_depth += 1 + if block_start is None: + block_start = i + 1 + elif ch == "}": + brace_depth -= 1 + if brace_depth == 0: + block_end = i + break + + match_body = raw_body[block_start:block_end] - # last arm - arms.append("".join(buf)) - arms = [arm for arm in arms if arm] + arm_pattern = re.compile( + r"((?:Self::\w+\s*\|\s*)*Self::\w+)\s*=>\s*(?:\{\s*)?Opcode::(\w+)", re.DOTALL + ) + variants_pattern = re.compile(r"Self::(\w+)") deopts = {} - for arm in arms: - *specialized, deopt = map(to_snake_case, re.findall(r"Self::(\w*)\b", arm)) - deopts[deopt] = specialized + for hit in arm_pattern.finditer(match_body): + raw_variants = hit.group(1) + opcode = hit.group(2) + + variants = variants_pattern.findall(raw_variants) + + key = to_snake_case(opcode) + value = [to_snake_case(variant) for variant in variants] + deopts[key] = value return deopts