gh-117139: Convert the evaluation stack to stack refs (#118450)

This PR sets up tagged pointers for CPython.

The general idea is to create a separate struct _PyStackRef for everything on the evaluation stack to store the bits. This forces the C compiler to warn us if we try to cast things or pull things out of the struct directly.

Only for free threading: We tag the low bit if something is deferred - that means we skip incref and decref operations on it. This behavior may change in the future if Mark's plans to defer all objects in the interpreter loop pans out.

This implies a strict stack reference discipline is required. ALL incref and decref operations on stackrefs must use the stackref variants. It is unsafe to untag something then do normal incref/decref ops on it.

The new incref and decref variants are called dup and close. They mimic a "handle" API operating on these stackrefs.

Please read Include/internal/pycore_stackref.h for more information!

---------

Co-authored-by: Mark Shannon <9448417+markshannon@users.noreply.github.com>

10 files changed, 56 insertions, 25 deletions

diff --git a/Tools/cases_generator/analyzer.py b/Tools/cases_generator/analyzer.py
index 96e2fd5..6b1af1b 100644
--- a/Tools/cases_generator/analyzer.py
+++ b/Tools/cases_generator/analyzer.py
@@ -111,7 +111,7 @@ class StackItem:
         return f"{type}{self.name}{size}{cond} {self.peek}"
 
     def is_array(self) -> bool:
-        return self.type == "PyObject **"
+        return self.type == "_PyStackRef *"
 
 
 @dataclass
@@ -353,6 +353,21 @@ def has_error_without_pop(op: parser.InstDef) -> bool:
 
 
 NON_ESCAPING_FUNCTIONS = (
+    "PyStackRef_FromPyObjectSteal",
+    "PyStackRef_AsPyObjectBorrow",
+    "PyStackRef_AsPyObjectSteal",
+    "PyStackRef_CLOSE",
+    "PyStackRef_DUP",
+    "PyStackRef_CLEAR",
+    "PyStackRef_IsNull",
+    "PyStackRef_TYPE",
+    "PyStackRef_False",
+    "PyStackRef_True",
+    "PyStackRef_None",
+    "PyStackRef_Is",
+    "PyStackRef_FromPyObjectNew",
+    "PyStackRef_AsPyObjectNew",
+    "PyStackRef_FromPyObjectImmortal",
     "Py_INCREF",
     "_PyManagedDictPointer_IsValues",
     "_PyObject_GetManagedDict",
@@ -399,8 +414,6 @@ NON_ESCAPING_FUNCTIONS = (
     "_PyFrame_SetStackPointer",
     "_PyType_HasFeature",
     "PyUnicode_Concat",
-    "_PyList_FromArraySteal",
-    "_PyTuple_FromArraySteal",
     "PySlice_New",
     "_Py_LeaveRecursiveCallPy",
     "CALL_STAT_INC",
@@ -413,6 +426,11 @@ NON_ESCAPING_FUNCTIONS = (
     "PyFloat_AS_DOUBLE",
     "_PyFrame_PushUnchecked",
     "Py_FatalError",
+    "STACKREFS_TO_PYOBJECTS",
+    "STACKREFS_TO_PYOBJECTS_CLEANUP",
+    "CONVERSION_FAILED",
+    "_PyList_FromArraySteal",
+    "_PyTuple_FromArraySteal",
 )
 
 ESCAPING_FUNCTIONS = (
diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py
index efbfc94..e4e0c9b 100644
--- a/Tools/cases_generator/generators_common.py
+++ b/Tools/cases_generator/generators_common.py
@@ -128,15 +128,15 @@ def replace_decrefs(
             continue
         if var.size != "1":
             out.emit(f"for (int _i = {var.size}; --_i >= 0;) {{\n")
-            out.emit(f"Py_DECREF({var.name}[_i]);\n")
+            out.emit(f"PyStackRef_CLOSE({var.name}[_i]);\n")
             out.emit("}\n")
         elif var.condition:
             if var.condition == "1":
-                out.emit(f"Py_DECREF({var.name});\n")
+                out.emit(f"PyStackRef_CLOSE({var.name});\n")
             elif var.condition != "0":
-                out.emit(f"Py_XDECREF({var.name});\n")
+                out.emit(f"PyStackRef_XCLOSE({var.name});\n")
         else:
-            out.emit(f"Py_DECREF({var.name});\n")
+            out.emit(f"PyStackRef_CLOSE({var.name});\n")
 
 
 def replace_sync_sp(
diff --git a/Tools/cases_generator/optimizer_generator.py b/Tools/cases_generator/optimizer_generator.py
index fb3e577..2775216 100644
--- a/Tools/cases_generator/optimizer_generator.py
+++ b/Tools/cases_generator/optimizer_generator.py
@@ -103,7 +103,7 @@ def write_uop(
         is_override = override is not None
         out.start_line()
         for var in reversed(prototype.stack.inputs):
-            res = stack.pop(var)
+            res = stack.pop(var, extract_bits=True)
             if not skip_inputs:
                 out.emit(res)
         if not prototype.properties.stores_sp:
@@ -140,7 +140,7 @@ def write_uop(
                 if not var.peek or is_override:
                     out.emit(stack.push(var))
         out.start_line()
-        stack.flush(out, cast_type="_Py_UopsSymbol *")
+        stack.flush(out, cast_type="_Py_UopsSymbol *", extract_bits=True)
     except SizeMismatch as ex:
         raise analysis_error(ex.args[0], uop.body[0])
 
diff --git a/Tools/cases_generator/parsing.py b/Tools/cases_generator/parsing.py
index cc897ff..0bd4229 100644
--- a/Tools/cases_generator/parsing.py
+++ b/Tools/cases_generator/parsing.py
@@ -285,7 +285,7 @@ class Parser(PLexer):
                 if not (size := self.expression()):
                     raise self.make_syntax_error("Expected expression")
                 self.require(lx.RBRACKET)
-                type_text = "PyObject **"
+                type_text = "_PyStackRef *"
                 size_text = size.text.strip()
             return StackEffect(tkn.text, type_text, cond_text, size_text)
         return None
diff --git a/Tools/cases_generator/stack.py b/Tools/cases_generator/stack.py
index 7325f2f..c0e1278 100644
--- a/Tools/cases_generator/stack.py
+++ b/Tools/cases_generator/stack.py
@@ -125,7 +125,7 @@ class Stack:
         self.variables: list[StackItem] = []
         self.defined: set[str] = set()
 
-    def pop(self, var: StackItem) -> str:
+    def pop(self, var: StackItem, extract_bits: bool = False) -> str:
         self.top_offset.pop(var)
         if not var.peek:
             self.peek_offset.pop(var)
@@ -155,8 +155,9 @@ class Stack:
         else:
             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 ""
         assign = (
-            f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}];"
+            f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}]{bits};"
         )
         if var.condition:
             if var.condition == "1":
@@ -178,11 +179,12 @@ class Stack:
             self.top_offset.push(var)
             return ""
 
-    def flush(self, out: CWriter, cast_type: str = "PyObject *") -> None:
+    def flush(self, out: CWriter, 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 ""
+                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":
@@ -190,7 +192,7 @@ class Stack:
                         elif var.condition != "1":
                             out.emit(f"if ({var.condition}) ")
                     out.emit(
-                        f"stack_pointer[{self.base_offset.to_c()}] = {cast}{var.name};\n"
+                        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():
diff --git a/Tools/cases_generator/tier1_generator.py b/Tools/cases_generator/tier1_generator.py
index 5df4413..c9dce1d 100644
--- a/Tools/cases_generator/tier1_generator.py
+++ b/Tools/cases_generator/tier1_generator.py
@@ -37,20 +37,25 @@ def declare_variables(inst: Instruction, out: CWriter) -> None:
         if isinstance(uop, Uop):
             for var in reversed(uop.stack.inputs):
                 if var.name not in variables:
-                    type = var.type if var.type else "PyObject *"
                     variables.add(var.name)
+                    type, null = (var.type, "NULL") if var.type else ("_PyStackRef", "PyStackRef_NULL")
+                    space = " " if type[-1].isalnum() else ""
                     if var.condition:
-                        out.emit(f"{type}{var.name} = NULL;\n")
+                        out.emit(f"{type}{space}{var.name} = {null};\n")
                     else:
-                        out.emit(f"{type}{var.name};\n")
+                        if var.is_array():
+                            out.emit(f"{var.type}{space}{var.name};\n")
+                        else:
+                            out.emit(f"{type}{space}{var.name};\n")
             for var in uop.stack.outputs:
                 if var.name not in variables:
                     variables.add(var.name)
-                    type = var.type if var.type else "PyObject *"
+                    type, null = (var.type, "NULL") if var.type else ("_PyStackRef", "PyStackRef_NULL")
+                    space = " " if type[-1].isalnum() else ""
                     if var.condition:
-                        out.emit(f"{type}{var.name} = NULL;\n")
+                        out.emit(f"{type}{space}{var.name} = {null};\n")
                     else:
-                        out.emit(f"{type}{var.name};\n")
+                        out.emit(f"{type}{space}{var.name};\n")
 
 
 def write_uop(
diff --git a/Tools/cases_generator/tier2_generator.py b/Tools/cases_generator/tier2_generator.py
index a091870..f3769bd 100644
--- a/Tools/cases_generator/tier2_generator.py
+++ b/Tools/cases_generator/tier2_generator.py
@@ -34,16 +34,17 @@ def declare_variable(
 ) -> None:
     if var.name in variables:
         return
-    type = var.type if var.type else "PyObject *"
     variables.add(var.name)
+    type, null = (var.type, "NULL") if var.type else ("_PyStackRef", "PyStackRef_NULL")
+    space = " " if type[-1].isalnum() else ""
     if var.condition:
-        out.emit(f"{type}{var.name} = NULL;\n")
+        out.emit(f"{type}{space}{var.name} = {null};\n")
         if uop.replicates:
             # Replicas may not use all their conditional variables
             # So avoid a compiler warning with a fake use
             out.emit(f"(void){var.name};\n")
     else:
-        out.emit(f"{type}{var.name};\n")
+        out.emit(f"{type}{space}{var.name};\n")
 
 
 def declare_variables(uop: Uop, out: CWriter) -> None:
diff --git a/Tools/gdb/libpython.py b/Tools/gdb/libpython.py
index 5fdc812..8aa7463 100755
--- a/Tools/gdb/libpython.py
+++ b/Tools/gdb/libpython.py
@@ -100,6 +100,8 @@ MAX_OUTPUT_LEN=1024
 
 hexdigits = "0123456789abcdef"
 
+USED_TAGS = 0b11
+
 ENCODING = locale.getpreferredencoding()
 
 FRAME_INFO_OPTIMIZED_OUT = '(frame information optimized out)'
@@ -158,6 +160,8 @@ class PyObjectPtr(object):
     _typename = 'PyObject'
 
     def __init__(self, gdbval, cast_to=None):
+        # Clear the tagged pointer
+        gdbval = gdb.Value(int(gdbval) & (~USED_TAGS)).cast(gdbval.type)
         if cast_to:
             self._gdbval = gdbval.cast(cast_to)
         else:
diff --git a/Tools/jit/template.c b/Tools/jit/template.c
index 39ce236..2bcbf8d 100644
--- a/Tools/jit/template.c
+++ b/Tools/jit/template.c
@@ -17,6 +17,7 @@
 #include "pycore_setobject.h"
 #include "pycore_sliceobject.h"
 #include "pycore_descrobject.h"
+#include "pycore_stackref.h"
 
 #include "ceval_macros.h"
 
@@ -84,7 +85,7 @@ do {                                                         \
 #define WITHIN_STACK_BOUNDS() 1
 
 _Py_CODEUNIT *
-_JIT_ENTRY(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *tstate)
+_JIT_ENTRY(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate)
 {
     // Locals that the instruction implementations expect to exist:
     PATCH_VALUE(_PyExecutorObject *, current_executor, _JIT_EXECUTOR)
diff --git a/Tools/jit/trampoline.c b/Tools/jit/trampoline.c
index 01b3d63..a0a963f 100644
--- a/Tools/jit/trampoline.c
+++ b/Tools/jit/trampoline.c
@@ -8,7 +8,7 @@
 // The actual change is patched in while the JIT compiler is being built, in
 // Tools/jit/_targets.py. On other platforms, this function compiles to nothing.
 _Py_CODEUNIT *
-_ENTRY(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *tstate)
+_ENTRY(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate)
 {
     // This is subtle. The actual trace will return to us once it exits, so we
     // need to make sure that we stay alive until then. If our trace side-exits