gh-139716: StackRef tests are added

Include/internal/pycore_stackref.h

@@ -74,6 +74,13 @@ static const _PyStackRef PyStackRef_ERROR = { .index = (1 << Py_TAGGED_SHIFT) };
 
 #define INITIAL_STACKREF_INDEX (5 << Py_TAGGED_SHIFT)
 
+static inline PyObject *
+_PyStackRef_AsTuple(_PyStackRef ref, PyObject *op)
+{
+    int flags = ref.index & Py_TAG_BITS;
+    return Py_BuildValue("(Ii)", Py_REFCNT(op), flags);
+}
+
 static inline _PyStackRef
 PyStackRef_Wrap(void *ptr)
 {
@@ -447,6 +454,13 @@ static const _PyStackRef PyStackRef_NULL = { .bits = Py_TAG_DEFERRED};
 
 #define PyStackRef_IsNullOrInt(stackref) (PyStackRef_IsNull(stackref) || PyStackRef_IsTaggedInt(stackref))
 
+static inline PyObject *
+_PyStackRef_AsTuple(_PyStackRef ref, PyObject *op)
+{
+    // Do not check StackRef flags in the free threading build.
+    return Py_BuildValue("(Ii)", Py_REFCNT(op), -1);
+}
+
 static inline PyObject *
 PyStackRef_AsPyObjectBorrow(_PyStackRef stackref)
 {
@@ -631,6 +645,13 @@ static const _PyStackRef PyStackRef_NULL = { .bits = PyStackRef_NULL_BITS };
 #define PyStackRef_IsFalse(REF) ((REF).bits == (((uintptr_t)&_Py_FalseStruct) | Py_TAG_REFCNT))
 #define PyStackRef_IsNone(REF) ((REF).bits == (((uintptr_t)&_Py_NoneStruct) | Py_TAG_REFCNT))
 
+static inline PyObject *
+_PyStackRef_AsTuple(_PyStackRef ref, PyObject *op)
+{
+    int flags = ref.bits & Py_TAG_BITS;
+    return Py_BuildValue("(Ii)", Py_REFCNT(op), flags);
+}
+
 #ifdef Py_DEBUG
 
 static inline void PyStackRef_CheckValid(_PyStackRef ref) {

Lib/test/test_stackrefs.py

@@ -0,0 +1,59 @@
+import unittest
+import sys
+from test.support import cpython_only
+try:
+    import _testinternalcapi
+except ImportError:
+    _testinternalcapi = None
+
+@cpython_only
+class TestDefinition(unittest.TestCase):
+
+    def test_equivalence(self):
+        def run_with_refcount_check(self, func, obj):
+            refcount = sys.getrefcount(obj)
+            res = func(obj)
+            self.assertEqual(sys.getrefcount(obj), refcount)
+            return res
+
+        funcs_with_incref = [
+            _testinternalcapi.stackref_from_object_new,
+            _testinternalcapi.stackref_from_object_steal_with_incref,
+            _testinternalcapi.stackref_make_heap_safe,
+            _testinternalcapi.stackref_make_heap_safe_with_borrow,
+            _testinternalcapi.stackref_strong_reference,
+        ]
+
+        funcs_with_borrow = [
+            _testinternalcapi.stackref_from_object_borrow,
+            _testinternalcapi.stackref_dup_borrowed_with_close,
+        ]
+
+        immortal_objs = (None, True, False, 42, '1')
+
+        for obj in immortal_objs:
+            results = set()
+            for func in funcs_with_incref + funcs_with_borrow:
+                res = run_with_refcount_check(self, func, obj)
+                results.add(res)
+            self.assertEqual(len(results), 1)
+
+        mortal_objs = (5000, 3+2j, range(10), object())
+
+        for obj in mortal_objs:
+            results = set()
+            for func in funcs_with_incref:
+                res = run_with_refcount_check(self, func, obj)
+                results.add(res)
+            self.assertEqual(len(results), 1)
+
+            results = set()
+            for func in funcs_with_borrow:
+                res = run_with_refcount_check(self, func, obj)
+                results.add(res)
+            self.assertEqual(len(results), 1)
+
+
+
+if __name__ == "__main__":
+    unittest.main()

Modules/_testinternalcapi.c

@@ -35,6 +35,7 @@
 #include "pycore_pylifecycle.h"   // _PyInterpreterConfig_InitFromDict()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_runtime_structs.h" // _PY_NSMALLPOSINTS
+#include "pycore_stackref.h"      // PyStackRef_FunctionCheck()
 #include "pycore_unicodeobject.h" // _PyUnicode_TransformDecimalAndSpaceToASCII()
 
 #include "clinic/_testinternalcapi.c.h"
@@ -2418,6 +2419,86 @@ set_vectorcall_nop(PyObject *self, PyObject *func)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+stackref_from_object_new(PyObject *self, PyObject *op)
+{
+    _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+    PyObject *obj = _PyStackRef_AsTuple(ref, op);
+    PyStackRef_CLOSE(ref);
+    return obj;
+}
+
+static PyObject *
+stackref_from_object_steal_with_incref(PyObject *self, PyObject *op)
+{
+    // The next two lines are equivalent to PyStackRef_FromPyObjectNew.
+    Py_INCREF(op);
+    _PyStackRef ref = PyStackRef_FromPyObjectSteal(op);
+    PyObject *obj = _PyStackRef_AsTuple(ref, op);
+    // The next two lines are equivalent to PyStackRef_CLOSE.
+    PyObject *op2 = PyStackRef_AsPyObjectSteal(ref);
+    Py_DECREF(op2);
+    return obj;
+}
+
+static PyObject *
+stackref_make_heap_safe(PyObject *self, PyObject *op)
+{
+    _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+    // For no-GIL release build ref2 is equal to ref.
+    _PyStackRef ref2 = PyStackRef_MakeHeapSafe(ref);
+    PyObject *obj = _PyStackRef_AsTuple(ref2, op);
+    PyStackRef_CLOSE(ref2);
+    return obj;
+}
+
+static PyObject *
+stackref_make_heap_safe_with_borrow(PyObject *self, PyObject *op)
+{
+    _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+    _PyStackRef ref2 = PyStackRef_Borrow(ref);
+    _PyStackRef ref3 = PyStackRef_MakeHeapSafe(ref2);
+    // We can close ref, since ref3 is heap safe.
+    PyStackRef_CLOSE(ref);
+    PyObject *obj = _PyStackRef_AsTuple(ref3, op);
+    PyStackRef_CLOSE(ref3);
+    return obj;
+}
+
+static PyObject *
+stackref_strong_reference(PyObject *self, PyObject *op)
+{
+    // The next two lines are equivalent to op2 = Py_NewRef(op).
+    _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+    PyObject *op2 = PyStackRef_AsPyObjectSteal(ref);
+    _PyStackRef ref2 = PyStackRef_FromPyObjectSteal(op2);
+    PyObject *obj = _PyStackRef_AsTuple(ref2, op);
+    PyStackRef_CLOSE(ref2);
+    return obj;
+}
+
+static PyObject *
+stackref_from_object_borrow(PyObject *self, PyObject *op)
+{
+    _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+    PyObject *obj = _PyStackRef_AsTuple(ref, op);
+    PyStackRef_CLOSE(ref);
+    return obj;
+}
+
+static PyObject *
+stackref_dup_borrowed_with_close(PyObject *self, PyObject *op)
+{
+    _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+    // For no-GIL release build ref2 is equal to ref.
+    _PyStackRef ref2 = PyStackRef_DUP(ref);
+    // For release build closing borrowed ref is no-op.
+    PyStackRef_XCLOSE(ref);
+    PyObject *obj = _PyStackRef_AsTuple(ref2, op);
+    PyStackRef_XCLOSE(ref2);
+    return obj;
+}
+
 static PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -2527,6 +2608,13 @@ static PyMethodDef module_functions[] = {
 #endif
     {"simple_pending_call", simple_pending_call, METH_O},
     {"set_vectorcall_nop", set_vectorcall_nop, METH_O},
+    {"stackref_from_object_new", stackref_from_object_new, METH_O},
+    {"stackref_from_object_steal_with_incref", stackref_from_object_steal_with_incref, METH_O},
+    {"stackref_make_heap_safe", stackref_make_heap_safe, METH_O},
+    {"stackref_make_heap_safe_with_borrow", stackref_make_heap_safe_with_borrow, METH_O},
+    {"stackref_strong_reference", stackref_strong_reference, METH_O},
+    {"stackref_from_object_borrow", stackref_from_object_borrow, METH_O},
+    {"stackref_dup_borrowed_with_close", stackref_dup_borrowed_with_close, METH_O},
     {NULL, NULL} /* sentinel */
 };