GH-122155: Track local variables between pops and pushes in cases generator (GH-122286)

Author: markshannon

Include/internal/pycore_opcode_metadata.h

@@ -31,7 +31,7 @@ def skip_if_different_mount_drives():
 with test_tools.imports_under_tool("cases_generator"):
     from analyzer import StackItem
     import parser
-    from stack import Stack
+    from stack import Local, Stack
     import tier1_generator
     import optimizer_generator
 
@@ -60,9 +60,9 @@ def test_effect_sizes(self):
         stack.pop(y)
         stack.pop(x)
         for out in outputs:
-            stack.push(out)
-        self.assertEqual(stack.base_offset.to_c(), "-1 - oparg*2 - oparg")
-        self.assertEqual(stack.top_offset.to_c(), "1 - oparg*2 - oparg + oparg*4")
+            stack.push(Local.local(out))
+        self.assertEqual(stack.base_offset.to_c(), "-1 - oparg - oparg*2")
+        self.assertEqual(stack.top_offset.to_c(), "1 - oparg - oparg*2 + oparg*4")
 
 
 class TestGeneratedCases(unittest.TestCase):
@@ -602,7 +602,11 @@ def test_array_error_if(self):
             frame->instr_ptr = next_instr;
             next_instr += 1;
             INSTRUCTION_STATS(OP);
-            if (oparg == 0) { stack_pointer += -1 - oparg; goto somewhere; }
+            if (oparg == 0) {
+                stack_pointer += -1 - oparg;
+                assert(WITHIN_STACK_BOUNDS());
+                goto somewhere;
+            }
             stack_pointer += -1 - oparg;
             assert(WITHIN_STACK_BOUNDS());
             DISPATCH();
@@ -908,19 +912,17 @@ def test_used_unused_used(self):
             next_instr += 1;
             INSTRUCTION_STATS(TEST);
             _PyStackRef w;
-            _PyStackRef x;
             _PyStackRef y;
             // FIRST
             w = stack_pointer[-1];
             {
                 use(w);
             }
             // SECOND
-            x = w;
             {
             }
             // THIRD
-            y = x;
+            y = w;
             {
                 use(y);
             }
@@ -1024,6 +1026,7 @@ def test_pop_on_error_peeks(self):
         }
 
         op(THIRD, (j, k --)) {
+            j,k; // Mark j and k as used
             ERROR_IF(cond, error);
         }
 
@@ -1054,6 +1057,7 @@ def test_pop_on_error_peeks(self):
             k = b;
             j = a;
             {
+                j,k; // Mark j and k as used
                 if (cond) goto pop_2_error;
             }
             stack_pointer += -2;
@@ -1063,6 +1067,51 @@ def test_pop_on_error_peeks(self):
         """
         self.run_cases_test(input, output)
 
+    def test_push_then_error(self):
+
+        input = """
+        op(FIRST, ( -- a)) {
+            a = 1;
+        }
+
+        op(SECOND, (a -- a, b)) {
+            b = 1;
+            ERROR_IF(cond, error);
+        }
+
+        macro(TEST) = FIRST + SECOND;
+        """
+
+        output = """
+        TARGET(TEST) {
+            frame->instr_ptr = next_instr;
+            next_instr += 1;
+            INSTRUCTION_STATS(TEST);
+            _PyStackRef a;
+            _PyStackRef b;
+            // FIRST
+            {
+                a = 1;
+            }
+            // SECOND
+            {
+                b = 1;
+                if (cond) {
+                    stack_pointer[0] = a;
+                    stack_pointer += 1;
+                    assert(WITHIN_STACK_BOUNDS());
+                    goto error;
+                }
+            }
+            stack_pointer[0] = a;
+            stack_pointer[1] = b;
+            stack_pointer += 2;
+            assert(WITHIN_STACK_BOUNDS());
+            DISPATCH();
+        }
+        """
+        self.run_cases_test(input, output)
+
 
 class TestGeneratedAbstractCases(unittest.TestCase):
     def setUp(self) -> None:

Lib/test/test_generated_cases.py

@@ -1357,8 +1357,8 @@ dummy_func(
             (void)counter;
         }
 
-        op(_UNPACK_SEQUENCE, (seq -- unused[oparg])) {
-            _PyStackRef *top = stack_pointer + oparg - 1;
+        op(_UNPACK_SEQUENCE, (seq -- output[oparg])) {
+            _PyStackRef *top = output + oparg;
             int res = _PyEval_UnpackIterableStackRef(tstate, seq, oparg, -1, top);
             DECREF_INPUTS();
             ERROR_IF(res == 0, error);
@@ -1401,9 +1401,8 @@ dummy_func(
             DECREF_INPUTS();
         }
 
-        inst(UNPACK_EX, (seq -- unused[oparg & 0xFF], unused, unused[oparg >> 8])) {
-            int totalargs = 1 + (oparg & 0xFF) + (oparg >> 8);
-            _PyStackRef *top = stack_pointer + totalargs - 1;
+        inst(UNPACK_EX, (seq -- left[oparg & 0xFF], unused, right[oparg >> 8])) {
+            _PyStackRef *top = right + (oparg >> 8);
             int res = _PyEval_UnpackIterableStackRef(tstate, seq, oparg & 0xFF, oparg >> 8, top);
             DECREF_INPUTS();
             ERROR_IF(res == 0, error);