gh-139716: StackRef tests are added

Include/internal/pycore_stackref.h

@@ -78,6 +78,13 @@ static const _PyStackRef PyStackRef_ERROR = { .index = (1 << Py_TAGGED_SHIFT) };
 
 #define PyStackRef_ZERO_BITS PyStackRef_NULL
 
+static inline PyObject *
+_PyStackRef_AsTuple(_PyStackRef ref, PyObject *op)
+{
+ int flags = ref.index & Py_TAG_BITS;
+ return Py_BuildValue("(ni)", Py_REFCNT(op), flags);
+}
+
 static inline _PyStackRef
 PyStackRef_Wrap(void *ptr)
 {
@@ -375,6 +382,19 @@ static const _PyStackRef PyStackRef_ERROR = { .bits = Py_TAG_INVALID };
  * PyStackRef_ZERO_BITS has no meaning and should not be used other than by the JIT. */
 static const _PyStackRef PyStackRef_ZERO_BITS = { .bits = 0 };
 
+static inline PyObject *
+_PyStackRef_AsTuple(_PyStackRef ref, PyObject *op)
+{
+ int flags = ref.bits & Py_TAG_BITS;
+#ifdef Py_GIL_DISABLED
+ if (_Py_IsImmortal(op)) {
+ // Do not check flags on immortal objects in the free threading build.
+ flags = -1;
+ }
+#endif
+ return Py_BuildValue("(ni)", Py_REFCNT(op), flags);
+}
+
 /* Wrap a pointer in a stack ref.
  * The resulting stack reference is not safe and should only be used
  * in the interpreter to pass values from one uop to another.

Lib/test/test_stackrefs.py

@@ -0,0 +1,71 @@
+import unittest
+import sys
+from test.support import cpython_only
+try:
+ import _testinternalcapi
+except ImportError:
+ _testinternalcapi = None
+
+@cpython_only
+class TestDefinition(unittest.TestCase):
+
+ BIG_REFCOUNT = 1000
+ FLAGS_REFCOUNT = 1
+ FLAGS_NO_REFCOUNT = 0
+ FLAGS_INVALID = -1
+
+ 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:
+ self.assertTrue(sys._is_immortal(obj))
+ results = set()
+ for func in funcs_with_incref + funcs_with_borrow:
+ refcount, flags = run_with_refcount_check(self, func, obj)
+ self.assertGreater(refcount, self.BIG_REFCOUNT)
+ self.assertIn(flags, (self.FLAGS_REFCOUNT, self.FLAGS_INVALID))
+ results.add((refcount, flags))
+ self.assertEqual(len(results), 1)
+
+ mortal_objs = (object(), range(10), iter([1, 2, 3]))
+
+ for obj in mortal_objs:
+ self.assertFalse(sys._is_immortal(obj))
+ results = set()
+ for func in funcs_with_incref:
+ refcount, flags = run_with_refcount_check(self, func, obj)
+ self.assertLess(refcount, self.BIG_REFCOUNT)
+ self.assertEqual(flags, self.FLAGS_NO_REFCOUNT)
+ results.add((refcount, flags))
+ self.assertEqual(len(results), 1)
+
+ results = set()
+ for func in funcs_with_borrow:
+ refcount, flags = run_with_refcount_check(self, func, obj)
+ self.assertLess(refcount, self.BIG_REFCOUNT)
+ self.assertEqual(flags, self.FLAGS_REFCOUNT)
+ results.add((refcount, flags))
+ 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"
@@ -2453,7 +2454,6 @@ module_get_gc_hooks(PyObject *self, PyObject *arg)
  return result;
 }
 
-
 static void
 check_threadstate_set_stack_protection(PyThreadState *tstate,
  void *start, size_t size)
@@ -2504,6 +2504,103 @@ test_threadstate_set_stack_protection(PyObject *self, PyObject *Py_UNUSED(args))
  Py_RETURN_NONE;
 }
 
+static PyObject *
+stackref_from_object_new(PyObject *self, PyObject *op)
+{
+ // Make a new strong reference and give ownership of it to ref.
+ _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+ PyObject *res = _PyStackRef_AsTuple(ref, op);
+ // Remove a strong reference by closing ref.
+ PyStackRef_CLOSE(ref);
+ return res;
+}
+
+static PyObject *
+stackref_from_object_steal_with_incref(PyObject *self, PyObject *op)
+{
+ // Combination of Py_INCREF and PyStackRef_FromPyObjectSteal
+ // is equivalent to PyStackRef_FromPyObjectNew.
+
+ Py_INCREF(op);
+ // Transfer ownership of a strong reference from op to ref.
+ _PyStackRef ref = PyStackRef_FromPyObjectSteal(op);
+
+ PyObject *res = _PyStackRef_AsTuple(ref, op);
+
+ // Combination of PyStackRef_AsPyObjectSteal and Py_DECREF
+ // is equivalent to PyStackRef_CLOSE.
+
+ // Transfer ownership of a strong reference from ref to op2.
+ PyObject *op2 = PyStackRef_AsPyObjectSteal(ref);
+ Py_DECREF(op2);
+
+ return res;
+}
+
+static PyObject *
+stackref_make_heap_safe(PyObject *self, PyObject *op)
+{
+ _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+ // Transfer ownership of a strong reference to op from ref to ref2,
+ // ref shouldn't be used anymore.
+ _PyStackRef ref2 = PyStackRef_MakeHeapSafe(ref);
+ PyObject *res = _PyStackRef_AsTuple(ref2, op);
+ PyStackRef_CLOSE(ref2);
+ return res;
+}
+
+static PyObject *
+stackref_make_heap_safe_with_borrow(PyObject *self, PyObject *op)
+{
+ _PyStackRef ref = PyStackRef_FromPyObjectNew(op);
+ // Create a borrowed reference to op, ref keeps a strong reference.
+ _PyStackRef ref2 = PyStackRef_Borrow(ref);
+ // Make a new strong reference to op from ref2,
+ // ref2 shouldn't be used anymore.
+ _PyStackRef ref3 = PyStackRef_MakeHeapSafe(ref2);
+ // We can close ref, since ref3 is heap safe.
+ PyStackRef_CLOSE(ref);
+ PyObject *res = _PyStackRef_AsTuple(ref3, op);
+ PyStackRef_CLOSE(ref3);
+ return res;
+}
+
+static PyObject *
+stackref_strong_reference(PyObject *self, PyObject *op)
+{
+ // Combination of PyStackRef_FromPyObjectBorrow and
+ // PyStackRef_AsPyObjectSteal is equivalent to Py_NewRef.
+ _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+ PyObject *op2 = PyStackRef_AsPyObjectSteal(ref);
+ _PyStackRef ref2 = PyStackRef_FromPyObjectSteal(op2);
+ PyObject *res = _PyStackRef_AsTuple(ref2, op);
+ PyStackRef_CLOSE(ref2);
+ return res;
+}
+
+static PyObject *
+stackref_from_object_borrow(PyObject *self, PyObject *op)
+{
+ _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+ PyObject *res = _PyStackRef_AsTuple(ref, op);
+ // Closing borrowed ref is no-op.
+ PyStackRef_CLOSE(ref);
+ return res;
+}
+
+static PyObject *
+stackref_dup_borrowed_with_close(PyObject *self, PyObject *op)
+{
+ _PyStackRef ref = PyStackRef_FromPyObjectBorrow(op);
+ // Duplicate reference, ref and ref2 both will be
+ // correct borrowed references from op.
+ _PyStackRef ref2 = PyStackRef_DUP(ref);
+ PyStackRef_XCLOSE(ref);
+ PyObject *res = _PyStackRef_AsTuple(ref2, op);
+ PyStackRef_XCLOSE(ref2);
+ return res;
+}
+
 
 static PyMethodDef module_functions[] = {
  {"get_configs", get_configs, METH_NOARGS},
@@ -2618,6 +2715,13 @@ static PyMethodDef module_functions[] = {
  {"module_get_gc_hooks", module_get_gc_hooks, METH_O},
  {"test_threadstate_set_stack_protection",
  test_threadstate_set_stack_protection, METH_NOARGS},
+ {"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 */
 };