python · markshannon · Jul 26, 2024 · Jul 23, 2024 · Jul 24, 2024
diff --git a/Lib/test/test_generated_cases.py b/Lib/test/test_generated_cases.py
@@ -813,6 +813,56 @@ def test_deopt_and_exit(self):
         with self.assertRaises(Exception):
             self.run_cases_test(input, output)
 
+    def test_pop_on_error_peeks(self):
+
+        input = """
+        op(FIRST, (x, y -- a, b)) {
+            a = x;
+            b = y;
+        }
+
+        op(SECOND, (a, b -- a, b)) {
+        }
+
+        op(THIRD, (j, k --)) {
+            ERROR_IF(cond, error);
+        }
+
+        macro(TEST) = FIRST + SECOND + THIRD;
+        """
+        output = """
+        TARGET(TEST) {
+            frame->instr_ptr = next_instr;
+            next_instr += 1;
+            INSTRUCTION_STATS(TEST);
+            PyObject *y;
+            PyObject *x;
+            PyObject *a;
+            PyObject *b;
+            PyObject *k;
+            PyObject *j;
+            // FIRST
+            y = stack_pointer[-1];
+            x = stack_pointer[-2];
+            {
+                a = x;
+                b = y;
+            }
+            // SECOND
+            {
+            }
+            // THIRD
+            k = b;
+            j = a;
+            {
+                if (cond) goto pop_2_error;
+            }
+            stack_pointer += -2;
+            DISPATCH();
+        }
+        """
+        self.run_cases_test(input, output)
+
 class TestGeneratedAbstractCases(unittest.TestCase):
     def setUp(self) -> None:
         super().setUp()

@@ -300,9 +300,13 @@ def analyze_stack(op: parser.InstDef, replace_op_arg_1: str | None = None) -> St
         convert_stack_item(i, replace_op_arg_1) for i in op.inputs if isinstance(i, parser.StackEffect)
     ]
     outputs: list[StackItem] = [convert_stack_item(i, replace_op_arg_1) for i in op.outputs]
+    # Mark variables with matching names at the base of the stack as "peek"
+    modified = False
     for input, output in zip(inputs, outputs):
-        if input.name == output.name:
+        if input.name == output.name and not modified:
             input.peek = output.peek = True
+        else:
+            modified = True
     return StackEffect(inputs, outputs)
 
 

@@ -84,7 +84,7 @@ def replace_error(
     next(tkn_iter)  # RPAREN
     next(tkn_iter)  # Semi colon
     out.emit(") ")
-    c_offset = stack.peek_offset.to_c()
+    c_offset = stack.peek_offset()
     try:
         offset = -int(c_offset)
         close = ";\n"

@@ -47,6 +47,9 @@ class StackOffset:
     def empty() -> "StackOffset":
         return StackOffset([], [])
 
+    def copy(self) -> "StackOffset":
+        return StackOffset(self.popped[:], self.pushed[:])
+
     def pop(self, item: StackItem) -> None:
         self.popped.append(var_size(item))
 
@@ -120,14 +123,11 @@ class Stack:
     def __init__(self) -> None:
         self.top_offset = StackOffset.empty()
         self.base_offset = StackOffset.empty()
-        self.peek_offset = StackOffset.empty()
         self.variables: list[StackItem] = []
         self.defined: set[str] = set()
 
     def pop(self, var: StackItem) -> str:
         self.top_offset.pop(var)
-        if not var.peek:
-            self.peek_offset.pop(var)
         indirect = "&" if var.is_array() else ""
         if self.variables:
             popped = self.variables.pop()
@@ -201,9 +201,16 @@ def flush(self, out: CWriter, cast_type: str = "PyObject *") -> None:
         self.variables = []
         self.base_offset.clear()
         self.top_offset.clear()
-        self.peek_offset.clear()
         out.start_line()
 
+    def peek_offset(self) -> str:
+        peek = self.base_offset.copy()
+        for var in self.variables:
+            if not var.peek:
+                break
+            peek.push(var)
+        return peek.to_c()
+
     def as_comment(self) -> str:
         return f"/* Variables: {[v.name for v in self.variables]}. Base offset: {self.base_offset.to_c()}. Top offset: {self.top_offset.to_c()} */"