gh-104615: don't make unsafe swaps in apply_static_swaps (#104620)

Author: carljm

Include/internal/pycore_compile.h

@@ -105,7 +105,8 @@ PyAPI_FUNC(PyObject*) _PyCompile_CodeGen(
 
 PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
         PyObject *instructions,
-        PyObject *consts);
+        PyObject *consts,
+        int nlocals);
 
 PyAPI_FUNC(PyCodeObject*)
 _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,

Include/internal/pycore_global_objects_fini_generated.h

@@ -559,6 +559,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(newline)
         STRUCT_FOR_ID(newlines)
         STRUCT_FOR_ID(next)
+        STRUCT_FOR_ID(nlocals)
         STRUCT_FOR_ID(node_depth)
         STRUCT_FOR_ID(node_offset)
         STRUCT_FOR_ID(ns)

Include/internal/pycore_global_strings.h

@@ -130,10 +130,10 @@ def generate_code(self, ast):
 
 class CfgOptimizationTestCase(CompilationStepTestCase):
 
-    def get_optimized(self, insts, consts):
+    def get_optimized(self, insts, consts, nlocals=0):
         insts = self.normalize_insts(insts)
         insts = self.complete_insts_info(insts)
-        insts = optimize_cfg(insts, consts)
+        insts = optimize_cfg(insts, consts, nlocals)
         return insts, consts
 
 class AssemblerTestCase(CompilationStepTestCase):

Include/internal/pycore_runtime_init_generated.h

@@ -1168,6 +1168,24 @@ def foo(param, lambda_exp):
         """)
         compile(code, "<test>", "exec")
 
+    def test_apply_static_swaps(self):
+        def f(x, y):
+            a, a = x, y
+            return a
+        self.assertEqual(f("x", "y"), "y")
+
+    def test_apply_static_swaps_2(self):
+        def f(x, y, z):
+            a, b, a = x, y, z
+            return a
+        self.assertEqual(f("x", "y", "z"), "z")
+
+    def test_apply_static_swaps_3(self):
+        def f(x, y, z):
+            a, a, b = x, y, z
+            return a
+        self.assertEqual(f("x", "y", "z"), "y")
+
 
 @requires_debug_ranges()
 class TestSourcePositions(unittest.TestCase):

Include/internal/pycore_unicodeobject_generated.h

@@ -484,6 +484,12 @@ def test_nested_listcomp_in_lambda(self):
         """
         self._check_in_scopes(code, {"z": 1, "out": [(3, 2, 1)]})
 
+    def test_assign_to_comp_iter_var_in_outer_function(self):
+        code = """
+            a = [1 for a in [0]]
+        """
+        self._check_in_scopes(code, {"a": [1]}, scopes=["function"])
+
 
 __test__ = {'doctests' : doctests}
 

Lib/test/support/bytecode_helper.py

@@ -971,13 +971,14 @@ def trace(frame, event, arg):
         self.assertNotInBytecode(f, "LOAD_FAST_CHECK")
 
 
-class DirectiCfgOptimizerTests(CfgOptimizationTestCase):
+class DirectCfgOptimizerTests(CfgOptimizationTestCase):
 
     def cfg_optimization_test(self, insts, expected_insts,
-                              consts=None, expected_consts=None):
+                              consts=None, expected_consts=None,
+                              nlocals=0):
         if expected_consts is None:
             expected_consts = consts
-        opt_insts, opt_consts = self.get_optimized(insts, consts)
+        opt_insts, opt_consts = self.get_optimized(insts, consts, nlocals)
         expected_insts = self.normalize_insts(expected_insts)
         self.assertInstructionsMatch(opt_insts, expected_insts)
         self.assertEqual(opt_consts, expected_consts)
@@ -1058,6 +1059,19 @@ def test_conditional_jump_backward_const_condition(self):
         ]
         self.cfg_optimization_test(insts, expected_insts, consts=list(range(5)))
 
+    def test_no_unsafe_static_swap(self):
+        # We can't change order of two stores to the same location
+        insts = [
+            ('LOAD_CONST', 0, 1),
+            ('LOAD_CONST', 1, 2),
+            ('LOAD_CONST', 2, 3),
+            ('SWAP', 3, 4),
+            ('STORE_FAST', 1, 4),
+            ('STORE_FAST', 1, 4),
+            ('POP_TOP', 0, 4),
+            ('RETURN_VALUE', 5)
+        ]
+        self.cfg_optimization_test(insts, insts, consts=list(range(3)), nlocals=1)
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_compile.py

@@ -0,0 +1,2 @@
+Fix wrong ordering of assignments in code like ``a, a = x, y``. Contributed by
+Carl Meyer.