GH-111485: Factor out tier 2 code generation from the rest of the interpreter code generator (GH-112968)

1 files changed, 87 insertions, 0 deletions

diff --git a/Tools/cases_generator/stack.py b/Tools/cases_generator/stack.py
index 9cd8e1a3b2e..c36a56ebf2d 100644
--- a/Tools/cases_generator/stack.py
+++ b/Tools/cases_generator/stack.py
@@ -2,6 +2,7 @@ import sys
 from analyzer import StackItem
 from dataclasses import dataclass
 from formatting import maybe_parenthesize
+from cwriter import CWriter
 
 
 def var_size(var: StackItem) -> str:
@@ -79,3 +80,89 @@ class StackOffset:
     def clear(self) -> None:
         self.popped = []
         self.pushed = []
+
+
+class SizeMismatch(Exception):
+    pass
+
+
+class Stack:
+    def __init__(self) -> None:
+        self.top_offset = StackOffset()
+        self.base_offset = StackOffset()
+        self.peek_offset = StackOffset()
+        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()
+            if popped.size != var.size:
+                raise SizeMismatch(
+                    f"Size mismatch when popping '{popped.name}' from stack to assign to {var.name}. "
+                    f"Expected {var.size} got {popped.size}"
+                )
+            if popped.name == var.name:
+                return ""
+            elif popped.name == "unused":
+                self.defined.add(var.name)
+                return (
+                    f"{var.name} = {indirect}stack_pointer[{self.top_offset.to_c()}];\n"
+                )
+            elif var.name == "unused":
+                return ""
+            else:
+                self.defined.add(var.name)
+                return f"{var.name} = {popped.name};\n"
+        self.base_offset.pop(var)
+        if var.name == "unused":
+            return ""
+        else:
+            self.defined.add(var.name)
+        cast = f"({var.type})" if (not indirect and var.type) else ""
+        assign = (
+            f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}];"
+        )
+        if var.condition:
+            return f"if ({var.condition}) {{ {assign} }}\n"
+        return f"{assign}\n"
+
+    def push(self, var: StackItem) -> str:
+        self.variables.append(var)
+        if var.is_array() and var.name not in self.defined and var.name != "unused":
+            c_offset = self.top_offset.to_c()
+            self.top_offset.push(var)
+            self.defined.add(var.name)
+            return f"{var.name} = &stack_pointer[{c_offset}];\n"
+        else:
+            self.top_offset.push(var)
+            return ""
+
+    def flush(self, out: CWriter) -> None:
+        for var in self.variables:
+            if not var.peek:
+                cast = "(PyObject *)" if var.type else ""
+                if var.name != "unused" and not var.is_array():
+                    if var.condition:
+                        out.emit(f" if ({var.condition}) ")
+                    out.emit(
+                        f"stack_pointer[{self.base_offset.to_c()}] = {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())
+            assert False
+        number = self.base_offset.to_c()
+        if number != "0":
+            out.emit(f"stack_pointer += {number};\n")
+        self.variables = []
+        self.base_offset.clear()
+        self.top_offset.clear()
+        self.peek_offset.clear()
+
+    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()} */"