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

7 files changed, 213 insertions, 102 deletions

diff --git a/Tools/cases_generator/analyzer.py b/Tools/cases_generator/analyzer.py
index 675dc0b..f6625a3 100644
--- a/Tools/cases_generator/analyzer.py
+++ b/Tools/cases_generator/analyzer.py
@@ -216,6 +216,7 @@ Part = Uop | Skip | Flush
 
 @dataclass
 class Instruction:
+    where: lexer.Token
     name: str
     parts: list[Part]
     _properties: Properties | None
@@ -690,9 +691,10 @@ def add_op(op: parser.InstDef, uops: dict[str, Uop]) -> None:
 
 
 def add_instruction(
-    name: str, parts: list[Part], instructions: dict[str, Instruction]
+    where: lexer.Token, name: str, parts: list[Part],
+    instructions: dict[str, Instruction]
 ) -> None:
-    instructions[name] = Instruction(name, parts, None)
+    instructions[name] = Instruction(where, name, parts, None)
 
 
 def desugar_inst(
@@ -720,7 +722,7 @@ def desugar_inst(
         parts.append(uop)
     else:
         parts[uop_index] = uop
-    add_instruction(name, parts, instructions)
+    add_instruction(inst.first_token, name, parts, instructions)
 
 
 def add_macro(
@@ -741,7 +743,7 @@ def add_macro(
             case _:
                 assert False
     assert parts
-    add_instruction(macro.name, parts, instructions)
+    add_instruction(macro.first_token, macro.name, parts, instructions)
 
 
 def add_family(
diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py
index 587dc0d..ab8c99f 100644
--- a/Tools/cases_generator/generators_common.py
+++ b/Tools/cases_generator/generators_common.py
@@ -8,7 +8,7 @@ from analyzer import (
     StackItem,
 )
 from cwriter import CWriter
-from typing import Callable, Mapping, TextIO, Iterator, Tuple
+from typing import Callable, Mapping, TextIO, Iterator
 from lexer import Token
 from stack import Stack
 
@@ -25,7 +25,7 @@ def root_relative_path(filename: str) -> str:
         return filename
 
 
-def type_and_null(var: StackItem) -> Tuple[str, str]:
+def type_and_null(var: StackItem) -> tuple[str, str]:
     if var.type:
         return var.type, "NULL"
     elif var.is_array():
@@ -95,16 +95,23 @@ def replace_error(
     c_offset = stack.peek_offset()
     try:
         offset = -int(c_offset)
-        close = ";\n"
     except ValueError:
-        offset = None
-        out.emit(f"{{ stack_pointer += {c_offset}; ")
-        close = "; }\n"
-    out.emit("goto ")
-    if offset:
-        out.emit(f"pop_{offset}_")
-    out.emit(label)
-    out.emit(close)
+        offset = -1
+    if offset > 0:
+        out.emit(f"goto pop_{offset}_")
+        out.emit(label)
+        out.emit(";\n")
+    elif offset == 0:
+        out.emit("goto ")
+        out.emit(label)
+        out.emit(";\n")
+    else:
+        out.emit("{\n")
+        stack.flush_locally(out)
+        out.emit("goto ")
+        out.emit(label)
+        out.emit(";\n")
+        out.emit("}\n")
 
 
 def replace_error_no_pop(
diff --git a/Tools/cases_generator/optimizer_generator.py b/Tools/cases_generator/optimizer_generator.py
index 6a66693..f6c2fea 100644
--- a/Tools/cases_generator/optimizer_generator.py
+++ b/Tools/cases_generator/optimizer_generator.py
@@ -23,7 +23,7 @@ from generators_common import (
 from cwriter import CWriter
 from typing import TextIO, Iterator
 from lexer import Token
-from stack import Stack, StackError
+from stack import Local, Stack, StackError
 
 DEFAULT_OUTPUT = ROOT / "Python/optimizer_cases.c.h"
 DEFAULT_ABSTRACT_INPUT = (ROOT / "Python/optimizer_bytecodes.c").absolute().as_posix()
@@ -98,19 +98,18 @@ def write_uop(
     debug: bool,
     skip_inputs: bool,
 ) -> None:
+    locals: dict[str, Local] = {}
     try:
         prototype = override if override else uop
         is_override = override is not None
         out.start_line()
         for var in reversed(prototype.stack.inputs):
-            res = stack.pop(var, extract_bits=True)
+            code, local = stack.pop(var, extract_bits=True)
             if not skip_inputs:
-                out.emit(res)
-        if not prototype.properties.stores_sp:
-            for i, var in enumerate(prototype.stack.outputs):
-                res = stack.push(var)
-                if not var.peek or is_override:
-                    out.emit(res)
+                out.emit(code)
+            if local.defined:
+                locals[local.name] = local
+        out.emit(stack.define_output_arrays(prototype.stack.outputs))
         if debug:
             args = []
             for var in prototype.stack.inputs:
@@ -135,10 +134,12 @@ def write_uop(
         else:
             emit_default(out, uop)
 
-        if prototype.properties.stores_sp:
-            for i, var in enumerate(prototype.stack.outputs):
-                if not var.peek or is_override:
-                    out.emit(stack.push(var))
+        for var in prototype.stack.outputs:
+            if var.name in locals:
+                local = locals[var.name]
+            else:
+                local = Local.local(var)
+            out.emit(stack.push(local))
         out.start_line()
         stack.flush(out, cast_type="_Py_UopsSymbol *", extract_bits=True)
     except StackError as ex:
diff --git a/Tools/cases_generator/parsing.py b/Tools/cases_generator/parsing.py
index 8957838..5acad57 100644
--- a/Tools/cases_generator/parsing.py
+++ b/Tools/cases_generator/parsing.py
@@ -60,6 +60,11 @@ class Node:
         end = context.end
         return tokens[begin:end]
 
+    @property
+    def first_token(self) -> lx.Token:
+        context = self.context
+        assert context is not None
+        return context.owner.tokens[context.begin]
 
 @dataclass
 class Block(Node):
diff --git a/Tools/cases_generator/stack.py b/Tools/cases_generator/stack.py
index 61dcfd3..d2d598a 100644
--- a/Tools/cases_generator/stack.py
+++ b/Tools/cases_generator/stack.py
@@ -39,6 +39,43 @@ def var_size(var: StackItem) -> str:
         return "1"
 
 @dataclass
+class Local:
+
+    item: StackItem
+    cached: bool
+    in_memory: bool
+    defined: bool
+
+    @staticmethod
+    def unused(defn: StackItem) -> "Local":
+        return Local(defn, False, defn.is_array(), False)
+
+    @staticmethod
+    def local(defn: StackItem) -> "Local":
+        array = defn.is_array()
+        return Local(defn, not array, array, True)
+
+    @staticmethod
+    def redefinition(var: StackItem, prev: "Local") -> "Local":
+        assert var.is_array() == prev.is_array()
+        return Local(var, prev.cached, prev.in_memory, True)
+
+    @property
+    def size(self) -> str:
+        return self.item.size
+
+    @property
+    def name(self) -> str:
+        return self.item.name
+
+    @property
+    def condition(self) -> str | None:
+        return self.item.condition
+
+    def is_array(self) -> bool:
+        return self.item.is_array()
+
+@dataclass
 class StackOffset:
     "The stack offset of the virtual base of the stack from the physical stack pointer"
 
@@ -66,7 +103,11 @@ class StackOffset:
 
     def simplify(self) -> None:
         "Remove matching values from both the popped and pushed list"
-        if not self.popped or not self.pushed:
+        if not self.popped:
+            self.pushed.sort()
+            return
+        if not self.pushed:
+            self.popped.sort()
             return
         # Sort the list so the lexically largest element is last.
         popped = sorted(self.popped)
@@ -87,6 +128,8 @@ class StackOffset:
                 popped.append(pop)
         self.popped.extend(popped)
         self.pushed.extend(pushed)
+        self.pushed.sort()
+        self.popped.sort()
 
     def to_c(self) -> str:
         self.simplify()
@@ -125,10 +168,10 @@ class Stack:
     def __init__(self) -> None:
         self.top_offset = StackOffset.empty()
         self.base_offset = StackOffset.empty()
-        self.variables: list[StackItem] = []
+        self.variables: list[Local] = []
         self.defined: set[str] = set()
 
-    def pop(self, var: StackItem, extract_bits: bool = False) -> str:
+    def pop(self, var: StackItem, extract_bits: bool = False) -> tuple[str, Local]:
         self.top_offset.pop(var)
         indirect = "&" if var.is_array() else ""
         if self.variables:
@@ -141,21 +184,32 @@ class Stack:
             if var.name in UNUSED:
                 if popped.name not in UNUSED and popped.name in self.defined:
                     raise StackError(f"Value is declared unused, but is already cached by prior operation")
-                return ""
-            if popped.name in UNUSED or popped.name not in self.defined:
-                self.defined.add(var.name)
+                return "", popped
+            if not var.used:
+                return "", popped
+            self.defined.add(var.name)
+            # Always define array variables as it is free, and their offset might have changed
+            if var.is_array():
+                return (
+                    f"{var.name} = &stack_pointer[{self.top_offset.to_c()}];\n",
+                    Local.redefinition(var, popped)
+                )
+            if not popped.defined:
                 return (
-                    f"{var.name} = {indirect}stack_pointer[{self.top_offset.to_c()}];\n"
+                    f"{var.name} = stack_pointer[{self.top_offset.to_c()}];\n",
+                    Local.redefinition(var, popped)
                 )
             else:
-                self.defined.add(var.name)
                 if popped.name == var.name:
-                    return ""
+                    return "", popped
                 else:
-                    return f"{var.name} = {popped.name};\n"
+                    return (
+                        f"{var.name} = {popped.name};\n",
+                        Local.redefinition(var, popped)
+                    )
         self.base_offset.pop(var)
         if var.name in UNUSED or not var.used:
-            return ""
+            return "", Local.unused(var)
         self.defined.add(var.name)
         cast = f"({var.type})" if (not indirect and var.type) else ""
         bits = ".bits" if cast and not extract_bits else ""
@@ -164,61 +218,80 @@ class Stack:
         )
         if var.condition:
             if var.condition == "1":
-                return f"{assign}\n"
+                assign = f"{assign}\n"
             elif var.condition == "0":
-                return ""
+               return "", Local.unused(var)
             else:
-                return f"if ({var.condition}) {{ {assign} }}\n"
-        return f"{assign}\n"
+                assign = f"if ({var.condition}) {{ {assign} }}\n"
+        else:
+            assign = f"{assign}\n"
+        in_memory = var.is_array() or var.peek
+        return assign, Local(var, not var.is_array(), in_memory, True)
 
-    def push(self, var: StackItem) -> str:
+    def push(self, var: Local) -> str:
         self.variables.append(var)
-        if var.is_array() and var.name not in self.defined and var.name not in UNUSED:
+        if var.is_array() and not var.defined and var.item.used:
+            assert var.in_memory
+            assert not var.cached
             c_offset = self.top_offset.to_c()
-            self.top_offset.push(var)
+            self.top_offset.push(var.item)
             self.defined.add(var.name)
+            var.defined = True
             return f"{var.name} = &stack_pointer[{c_offset}];\n"
         else:
-            self.top_offset.push(var)
-            if var.used:
+            self.top_offset.push(var.item)
+            if var.item.used:
                 self.defined.add(var.name)
             return ""
 
-    def flush(self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = False) -> None:
+    def define_output_arrays(self, outputs: list[StackItem]) -> str:
+        res = []
+        top_offset = self.top_offset.copy()
+        for var in outputs:
+            if var.is_array() and var.used and not var.peek:
+                c_offset = top_offset.to_c()
+                top_offset.push(var)
+                res.append(f"{var.name} = &stack_pointer[{c_offset}];\n")
+            else:
+                top_offset.push(var)
+        return "\n".join(res)
+
+    @staticmethod
+    def _do_flush(out: CWriter, variables: list[Local], base_offset: StackOffset, top_offset: StackOffset,
+                  cast_type: str = "uintptr_t", extract_bits: bool = False) -> None:
         out.start_line()
-        for var in self.variables:
-            if not var.peek:
-                cast = f"({cast_type})" if var.type else ""
+        for var in variables:
+            if var.cached and not var.in_memory and not var.item.peek and not var.name in UNUSED:
+                cast = f"({cast_type})" if var.item.type else ""
                 bits = ".bits" if cast and not extract_bits else ""
-                if var.name not in UNUSED and not var.is_array():
-                    if var.condition:
-                        if var.condition == "0":
-                            continue
-                        elif var.condition != "1":
-                            out.emit(f"if ({var.condition}) ")
-                    out.emit(
-                        f"stack_pointer[{self.base_offset.to_c()}]{bits} = {cast}{var.name};\n"
-                    )
-            self.base_offset.push(var)
-        if self.base_offset.to_c() != self.top_offset.to_c():
-            print("base", self.base_offset.to_c(), "top", self.top_offset.to_c())
+                if var.condition == "0":
+                    continue
+                if var.condition and var.condition != "1":
+                    out.emit(f"if ({var.condition}) ")
+                out.emit(
+                    f"stack_pointer[{base_offset.to_c()}]{bits} = {cast}{var.name};\n"
+                )
+            base_offset.push(var.item)
+        if base_offset.to_c() != top_offset.to_c():
+            print("base", base_offset, "top", top_offset)
             assert False
-        number = self.base_offset.to_c()
+        number = base_offset.to_c()
         if number != "0":
             out.emit(f"stack_pointer += {number};\n")
             out.emit("assert(WITHIN_STACK_BOUNDS());\n")
+        out.start_line()
+
+    def flush_locally(self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = False) -> None:
+        self._do_flush(out, self.variables[:], self.base_offset.copy(), self.top_offset.copy(), cast_type, extract_bits)
+
+    def flush(self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = False) -> None:
+        self._do_flush(out, self.variables, self.base_offset, self.top_offset, cast_type, extract_bits)
         self.variables = []
         self.base_offset.clear()
         self.top_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()
+        return self.top_offset.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()} */"
@@ -236,8 +309,15 @@ def get_stack_effect(inst: Instruction | PseudoInstruction) -> Stack:
             yield inst.stack
 
     for s in stacks(inst):
+        locals: dict[str, Local] = {}
         for var in reversed(s.inputs):
-            stack.pop(var)
+            _, local = stack.pop(var)
+            if var.name != "unused":
+                locals[local.name] = local
         for var in s.outputs:
-            stack.push(var)
+            if var.name in locals:
+                local = locals[var.name]
+            else:
+                local = Local.unused(var)
+            stack.push(local)
     return stack
diff --git a/Tools/cases_generator/tier1_generator.py b/Tools/cases_generator/tier1_generator.py
index 118f4b3..1cdafbd 100644
--- a/Tools/cases_generator/tier1_generator.py
+++ b/Tools/cases_generator/tier1_generator.py
@@ -25,7 +25,7 @@ from generators_common import (
 )
 from cwriter import CWriter
 from typing import TextIO
-from stack import Stack, StackError
+from stack import Local, Stack, StackError, get_stack_effect
 
 
 DEFAULT_OUTPUT = ROOT / "Python/generated_cases.c.h"
@@ -43,18 +43,12 @@ def declare_variable(var: StackItem, out: CWriter) -> None:
 
 
 def declare_variables(inst: Instruction, out: CWriter) -> None:
-    stack = Stack()
-    for part in inst.parts:
-        if not isinstance(part, Uop):
-            continue
-        try:
-            for var in reversed(part.stack.inputs):
-                stack.pop(var)
-            for var in part.stack.outputs:
-                 stack.push(var)
-        except StackError as ex:
-            raise analysis_error(ex.args[0], part.body[0]) from None
+    try:
+        stack = get_stack_effect(inst)
+    except StackError as ex:
+        raise analysis_error(ex.args[0], inst.where)
     required = set(stack.defined)
+    required.discard("unused")
     for part in inst.parts:
         if not isinstance(part, Uop):
             continue
@@ -80,16 +74,26 @@ def write_uop(
         stack.flush(out)
         return offset
     try:
+        locals: dict[str, Local] = {}
         out.start_line()
         if braces:
             out.emit(f"// {uop.name}\n")
+        peeks: list[Local] = []
         for var in reversed(uop.stack.inputs):
-            out.emit(stack.pop(var))
+            code, local = stack.pop(var)
+            out.emit(code)
+            if var.peek:
+                peeks.append(local)
+            if local.defined:
+                locals[local.name] = local
+        # Push back the peeks, so that they remain on the logical
+        # stack, but their values are cached.
+        while peeks:
+            stack.push(peeks.pop())
         if braces:
             out.emit("{\n")
-        if not uop.properties.stores_sp:
-            for i, var in enumerate(uop.stack.outputs):
-                out.emit(stack.push(var))
+        out.emit(stack.define_output_arrays(uop.stack.outputs))
+
         for cache in uop.caches:
             if cache.name != "unused":
                 if cache.size == 4:
@@ -105,16 +109,22 @@ def write_uop(
                     out.emit(f"(void){cache.name};\n")
             offset += cache.size
         emit_tokens(out, uop, stack, inst)
-        if uop.properties.stores_sp:
-            for i, var in enumerate(uop.stack.outputs):
-                out.emit(stack.push(var))
+        for i, var in enumerate(uop.stack.outputs):
+            if not var.peek:
+                if var.name in locals:
+                    local = locals[var.name]
+                elif var.name == "unused":
+                    local = Local.unused(var)
+                else:
+                    local = Local.local(var)
+                out.emit(stack.push(local))
         if braces:
             out.start_line()
             out.emit("}\n")
         # out.emit(stack.as_comment() + "\n")
         return offset
     except StackError as ex:
-        raise analysis_error(ex.args[0], uop.body[0]) from None
+        raise analysis_error(ex.args[0], uop.body[0])
 
 
 def uses_this(inst: Instruction) -> bool:
diff --git a/Tools/cases_generator/tier2_generator.py b/Tools/cases_generator/tier2_generator.py
index 88ad0fd..18bab2c 100644
--- a/Tools/cases_generator/tier2_generator.py
+++ b/Tools/cases_generator/tier2_generator.py
@@ -25,7 +25,7 @@ from generators_common import (
 from cwriter import CWriter
 from typing import TextIO, Iterator
 from lexer import Token
-from stack import Stack, StackError
+from stack import Local, Stack, StackError, get_stack_effect
 
 DEFAULT_OUTPUT = ROOT / "Python/executor_cases.c.h"
 
@@ -53,8 +53,9 @@ def declare_variables(uop: Uop, out: CWriter) -> None:
     for var in reversed(uop.stack.inputs):
         stack.pop(var)
     for var in uop.stack.outputs:
-            stack.push(var)
+        stack.push(Local.unused(var))
     required = set(stack.defined)
+    required.discard("unused")
     for var in reversed(uop.stack.inputs):
         declare_variable(var, uop, required, out)
     for var in uop.stack.outputs:
@@ -156,6 +157,7 @@ TIER2_REPLACEMENT_FUNCTIONS["EXIT_IF"] = tier2_replace_exit_if
 
 
 def write_uop(uop: Uop, out: CWriter, stack: Stack) -> None:
+    locals: dict[str, Local] = {}
     try:
         out.start_line()
         if uop.properties.oparg:
@@ -165,10 +167,11 @@ def write_uop(uop: Uop, out: CWriter, stack: Stack) -> None:
             out.emit(f"oparg = {uop.properties.const_oparg};\n")
             out.emit(f"assert(oparg == CURRENT_OPARG());\n")
         for var in reversed(uop.stack.inputs):
-            out.emit(stack.pop(var))
-        if not uop.properties.stores_sp:
-            for i, var in enumerate(uop.stack.outputs):
-                out.emit(stack.push(var))
+            code, local = stack.pop(var)
+            out.emit(code)
+            if local.defined:
+                locals[local.name] = local
+        out.emit(stack.define_output_arrays(uop.stack.outputs))
         for cache in uop.caches:
             if cache.name != "unused":
                 if cache.size == 4:
@@ -178,9 +181,12 @@ def write_uop(uop: Uop, out: CWriter, stack: Stack) -> None:
                     cast = f"uint{cache.size*16}_t"
                 out.emit(f"{type}{cache.name} = ({cast})CURRENT_OPERAND();\n")
         emit_tokens(out, uop, stack, None, TIER2_REPLACEMENT_FUNCTIONS)
-        if uop.properties.stores_sp:
-            for i, var in enumerate(uop.stack.outputs):
-                out.emit(stack.push(var))
+        for i, var in enumerate(uop.stack.outputs):
+            if var.name in locals:
+                local = locals[var.name]
+            else:
+                local = Local.local(var)
+            out.emit(stack.push(local))
     except StackError as ex:
         raise analysis_error(ex.args[0], uop.body[0]) from None