GH-122155: Fix cases generator to correctly compute 'peek' offset for error handling (GH-122158)

Lib/test/test_generated_cases.py

@@ -1012,6 +1012,57 @@ class TestGeneratedCases(unittest.TestCase):
         """
         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);
+            _PyStackRef x;
+            _PyStackRef y;
+            _PyStackRef a;
+            _PyStackRef b;
+            _PyStackRef j;
+            _PyStackRef k;
+            // 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;
+            assert(WITHIN_STACK_BOUNDS());
+            DISPATCH();
+        }
+        """
+        self.run_cases_test(input, output)
+
 
 class TestGeneratedAbstractCases(unittest.TestCase):
     def setUp(self) -> None:

Tools/cases_generator/analyzer.py

@@ -317,9 +317,13 @@ def analyze_stack(op: parser.InstDef | parser.Pseudo, replace_op_arg_1: str | No
         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
     if isinstance(op, parser.InstDef):
         output_names = [out.name for out in outputs]
         for input in inputs:

Tools/cases_generator/generators_common.py

@@ -92,7 +92,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"

Tools/cases_generator/stack.py

@@ -49,6 +49,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))
 
@@ -122,14 +125,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, extract_bits: bool = False) -> 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()
@@ -210,9 +210,16 @@ class Stack:
         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()} */"