GH-128682: Stronger checking of `PyStackRef_CLOSE` and `DEAD`. (GH-128683)

1 files changed, 88 insertions, 15 deletions

diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py
index 8df9a9c..0035deb 100644
--- a/Tools/cases_generator/generators_common.py
+++ b/Tools/cases_generator/generators_common.py
@@ -98,6 +98,11 @@ def always_true(tkn: Token | None) -> bool:
         return False
     return tkn.text in {"true", "1"}
 
+NON_ESCAPING_DEALLOCS = {
+    "_PyFloat_ExactDealloc",
+    "_PyLong_ExactDealloc",
+    "_PyUnicode_ExactDealloc",
+}
 
 class Emitter:
     out: CWriter
@@ -116,7 +121,7 @@ class Emitter:
             "SAVE_STACK": self.save_stack,
             "RELOAD_STACK": self.reload_stack,
             "PyStackRef_CLOSE": self.stackref_close,
-            "PyStackRef_CLOSE_SPECIALIZED": self.stackref_close,
+            "PyStackRef_CLOSE_SPECIALIZED": self.stackref_close_specialized,
             "PyStackRef_AsPyObjectSteal": self.stackref_steal,
             "DISPATCH": self.dispatch,
             "INSTRUCTION_SIZE": self.instruction_size,
@@ -234,23 +239,26 @@ class Emitter:
         next(tkn_iter)
         next(tkn_iter)
         self.out.emit_at("", tkn)
-        for var in uop.stack.inputs:
-            if var.name == "unused" or var.name == "null" or var.peek:
+        for var in storage.inputs:
+            if not var.defined:
+                continue
+            if var.name == "null":
                 continue
+            close = "PyStackRef_CLOSE"
+            if "null" in var.name or var.condition and var.condition != "1":
+                close = "PyStackRef_XCLOSE"
             if var.size:
                 if var.size == "1":
-                    self.out.emit(f"PyStackRef_CLOSE({var.name}[0]);\n")
+                    self.out.emit(f"{close}({var.name}[0]);\n")
                 else:
                     self.out.emit(f"for (int _i = {var.size}; --_i >= 0;) {{\n")
-                    self.out.emit(f"PyStackRef_CLOSE({var.name}[_i]);\n")
+                    self.out.emit(f"{close}({var.name}[_i]);\n")
                     self.out.emit("}\n")
             elif var.condition:
-                if var.condition == "1":
-                    self.out.emit(f"PyStackRef_CLOSE({var.name});\n")
-                elif var.condition != "0":
-                    self.out.emit(f"PyStackRef_XCLOSE({var.name});\n")
+                if var.condition != "0":
+                    self.out.emit(f"{close}({var.name});\n")
             else:
-                self.out.emit(f"PyStackRef_CLOSE({var.name});\n")
+                self.out.emit(f"{close}({var.name});\n")
         for input in storage.inputs:
             input.defined = False
         return True
@@ -291,6 +299,25 @@ class Emitter:
             raise analysis_error(f"'{name}' is not a live input-only variable", name_tkn)
         return True
 
+    def stackref_kill(
+        self,
+        name: Token,
+        storage: Storage,
+        escapes: bool
+    ) -> bool:
+        live = ""
+        for var in reversed(storage.inputs):
+            if var.name == name.text:
+                if live and escapes:
+                    raise analysis_error(
+                        f"Cannot close '{name.text}' when "
+                        f"'{live}' is still live", name)
+                var.defined = False
+                break
+            if var.defined:
+                live = var.name
+        return True
+
     def stackref_close(
         self,
         tkn: Token,
@@ -306,14 +333,60 @@ class Emitter:
         name = next(tkn_iter)
         self.out.emit(name)
         if name.kind == "IDENTIFIER":
-            for var in storage.inputs:
-                if var.name == name.text:
-                    var.defined = False
+            return self.stackref_kill(name, storage, True)
         rparen = emit_to(self.out, tkn_iter, "RPAREN")
         self.emit(rparen)
         return True
 
-    stackref_steal = stackref_close
+    def stackref_close_specialized(
+        self,
+        tkn: Token,
+        tkn_iter: TokenIterator,
+        uop: Uop,
+        storage: Storage,
+        inst: Instruction | None,
+    ) -> bool:
+
+        self.out.emit(tkn)
+        tkn = next(tkn_iter)
+        assert tkn.kind == "LPAREN"
+        self.out.emit(tkn)
+        name = next(tkn_iter)
+        self.out.emit(name)
+        comma = next(tkn_iter)
+        if comma.kind != "COMMA":
+            raise analysis_error("Expected comma", comma)
+        self.out.emit(comma)
+        dealloc = next(tkn_iter)
+        if dealloc.kind != "IDENTIFIER":
+             raise analysis_error("Expected identifier", dealloc)
+        self.out.emit(dealloc)
+        if name.kind == "IDENTIFIER":
+            escapes = dealloc.text not in NON_ESCAPING_DEALLOCS
+            return self.stackref_kill(name, storage, escapes)
+        rparen = emit_to(self.out, tkn_iter, "RPAREN")
+        self.emit(rparen)
+        return True
+
+    def stackref_steal(
+        self,
+        tkn: Token,
+        tkn_iter: TokenIterator,
+        uop: Uop,
+        storage: Storage,
+        inst: Instruction | None,
+    ) -> bool:
+        self.out.emit(tkn)
+        tkn = next(tkn_iter)
+        assert tkn.kind == "LPAREN"
+        self.out.emit(tkn)
+        name = next(tkn_iter)
+        self.out.emit(name)
+        if name.kind == "IDENTIFIER":
+            return self.stackref_kill(name, storage, False)
+        rparen = emit_to(self.out, tkn_iter, "RPAREN")
+        self.emit(rparen)
+        return True
 
     def sync_sp(
         self,
@@ -548,7 +621,7 @@ class Emitter:
             storage.push_outputs()
             self._print_storage(storage)
         except StackError as ex:
-            raise analysis_error(ex.args[0], rbrace)
+            raise analysis_error(ex.args[0], rbrace) from None
         return storage
 
     def emit(self, txt: str | Token) -> None: