diff --git a/Include/internal/pycore_stackref.h b/Include/internal/pycore_stackref.h index 94fcb1d8aee52b..dd04a3f78caaf1 100644 --- a/Include/internal/pycore_stackref.h +++ b/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) { diff --git a/Lib/test/test_stackrefs.py b/Lib/test/test_stackrefs.py new file mode 100644 index 00000000000000..955a5e7b5b0a45 --- /dev/null +++ b/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() diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c index c2647d405e25bc..d3f4158a40aac2 100644 --- a/Modules/_testinternalcapi.c +++ b/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 */ };