The same pattern of refcounting for stackrefs as in production build

Author: efimov-mikhail

Include/internal/pycore_stackref.h

@@ -53,36 +53,53 @@ extern "C" {
 
 #if !defined(Py_GIL_DISABLED) && defined(Py_STACKREF_DEBUG)
 
-#define Py_TAG_BITS 0
+#define Py_INT_TAG 3
+#define Py_TAG_INVALID 2
+#define Py_TAG_REFCNT 1
+#define Py_TAG_BITS 3
 
 PyAPI_FUNC(PyObject *) _Py_stackref_get_object(_PyStackRef ref);
 PyAPI_FUNC(PyObject *) _Py_stackref_close(_PyStackRef ref, const char *filename, int linenumber);
-PyAPI_FUNC(_PyStackRef) _Py_stackref_create(PyObject *obj, const char *filename, int linenumber);
+PyAPI_FUNC(_PyStackRef) _Py_stackref_create(PyObject *obj, uint16_t flags, const char *filename, int linenumber);
 PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber);
 extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);
 
 static const _PyStackRef PyStackRef_NULL = { .index = 0 };
-static const _PyStackRef PyStackRef_ERROR = { .index = 2 };
+static const _PyStackRef PyStackRef_ERROR = { .index = 4 };
 
-// Use the first 3 even numbers for None, True and False.
-// Odd numbers are reserved for (tagged) integers
-#define PyStackRef_None ((_PyStackRef){ .index = 4 } )
-#define PyStackRef_False ((_PyStackRef){ .index = 6 })
-#define PyStackRef_True ((_PyStackRef){ .index = 8 })
+#define PyStackRef_None ((_PyStackRef){ .index = 8 } )
+#define PyStackRef_False ((_PyStackRef){ .index = 12 })
+#define PyStackRef_True ((_PyStackRef){ .index = 16 })
 
-#define INITIAL_STACKREF_INDEX 10
+#define INITIAL_STACKREF_INDEX 20
 
 static inline _PyStackRef
 PyStackRef_Wrap(void *ptr)
 {
     assert(ptr != NULL);
+#ifdef Py_DEBUG
+    assert(((uint64_t)ptr & Py_TAG_BITS) == 0);
+    return (_PyStackRef){ .index = ((uint64_t)ptr) | Py_TAG_INVALID };
+#else
     return (_PyStackRef){ .index = (uint64_t)ptr };
+#endif
 }
 
 static inline void *
 PyStackRef_Unwrap(_PyStackRef ref)
 {
+#ifdef Py_DEBUG
+    assert ((ref.index & Py_TAG_BITS) == Py_TAG_INVALID);
+    return (void *)(ref.index & ~Py_TAG_BITS);
+#else
     return (void *)(ref.index);
+#endif
+}
+
+static inline int
+PyStackRef_RefcountOnObject(_PyStackRef ref)
+{
+    return (ref.index & Py_TAG_REFCNT) == 0;
 }
 
 static inline int
@@ -94,7 +111,7 @@ PyStackRef_IsNull(_PyStackRef ref)
 static inline bool
 PyStackRef_IsError(_PyStackRef ref)
 {
-    return ref.index == 2;
+    return ref.index == 4;
 }
 
 static inline bool
@@ -125,7 +142,7 @@ PyStackRef_IsNone(_PyStackRef ref)
 static inline bool
 PyStackRef_IsTaggedInt(_PyStackRef ref)
 {
-    return (ref.index & 1) == 1;
+    return (ref.index & Py_TAG_BITS) == Py_INT_TAG;
 }
 
 static inline PyObject *
@@ -136,51 +153,68 @@ _PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumb
     _Py_stackref_record_borrow(ref, filename, linenumber);
     return _Py_stackref_get_object(ref);
 }
-
 #define PyStackRef_AsPyObjectBorrow(REF) _PyStackRef_AsPyObjectBorrow((REF), __FILE__, __LINE__)
 
 static inline PyObject *
 _PyStackRef_AsPyObjectSteal(_PyStackRef ref, const char *filename, int linenumber)
 {
-    return _Py_stackref_close(ref, filename, linenumber);
+    PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
+    if (PyStackRef_RefcountOnObject(ref)) {
+        return obj;
+    }
+    return Py_NewRef(obj);
 }
 #define PyStackRef_AsPyObjectSteal(REF) _PyStackRef_AsPyObjectSteal((REF), __FILE__, __LINE__)
 
 static inline _PyStackRef
 _PyStackRef_FromPyObjectNew(PyObject *obj, const char *filename, int linenumber)
 {
-    Py_INCREF(obj);
-    return _Py_stackref_create(obj, filename, linenumber);
+    assert(obj != NULL);
+    uint16_t flags = 0;
+    if (!_Py_IsImmortal(obj)) {
+        _Py_INCREF_MORTAL(obj);
+    } else {
+        flags = Py_TAG_REFCNT;
+    }
+    return _Py_stackref_create(obj, flags, filename, linenumber);
 }
 #define PyStackRef_FromPyObjectNew(obj) _PyStackRef_FromPyObjectNew(_PyObject_CAST(obj), __FILE__, __LINE__)
 
 static inline _PyStackRef
 _PyStackRef_FromPyObjectSteal(PyObject *obj, const char *filename, int linenumber)
 {
-    return _Py_stackref_create(obj, filename, linenumber);
+    assert(obj != NULL);
+    uint16_t flags = 0;
+    if (_Py_IsImmortal(obj)) {
+        flags = Py_TAG_REFCNT;
+    }
+    return _Py_stackref_create(obj, flags, filename, linenumber);
 }
 #define PyStackRef_FromPyObjectSteal(obj) _PyStackRef_FromPyObjectSteal(_PyObject_CAST(obj), __FILE__, __LINE__)
 
 static inline _PyStackRef
 _PyStackRef_FromPyObjectBorrow(PyObject *obj, const char *filename, int linenumber)
 {
-    Py_INCREF(obj);
-    return _Py_stackref_create(obj, filename, linenumber);
+    return _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
 }
 #define PyStackRef_FromPyObjectBorrow(obj) _PyStackRef_FromPyObjectBorrow(_PyObject_CAST(obj), __FILE__, __LINE__)
 
 static inline void
 _PyStackRef_CLOSE(_PyStackRef ref, const char *filename, int linenumber)
 {
+    assert(!PyStackRef_IsError(ref));
+    assert(!PyStackRef_IsNull(ref));
     if (PyStackRef_IsTaggedInt(ref)) {
         return;
     }
     PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
-    Py_DECREF(obj);
+    assert(Py_REFCNT(obj) > 0);
+    if (PyStackRef_RefcountOnObject(ref)) {
+        Py_DECREF(obj);
+    }
 }
 #define PyStackRef_CLOSE(REF) _PyStackRef_CLOSE((REF), __FILE__, __LINE__)
 
-
 static inline void
 _PyStackRef_XCLOSE(_PyStackRef ref, const char *filename, int linenumber)
 {
@@ -196,31 +230,46 @@ static inline _PyStackRef
 _PyStackRef_DUP(_PyStackRef ref, const char *filename, int linenumber)
 {
     assert(!PyStackRef_IsError(ref));
+    assert(!PyStackRef_IsNull(ref));
     if (PyStackRef_IsTaggedInt(ref)) {
         return ref;
     }
-    else {
-        PyObject *obj = _Py_stackref_get_object(ref);
+    PyObject *obj = _Py_stackref_get_object(ref);
+    uint16_t flags = 0;
+    if (PyStackRef_RefcountOnObject(ref)) {
         Py_INCREF(obj);
-        return _Py_stackref_create(obj, filename, linenumber);
+    } else {
+        flags = Py_TAG_REFCNT;
     }
+    return _Py_stackref_create(obj, flags, filename, linenumber);
 }
 #define PyStackRef_DUP(REF) _PyStackRef_DUP(REF, __FILE__, __LINE__)
 
-extern void _PyStackRef_CLOSE_SPECIALIZED(_PyStackRef ref, destructor destruct, const char *filename, int linenumber);
-#define PyStackRef_CLOSE_SPECIALIZED(REF, DESTRUCT) _PyStackRef_CLOSE_SPECIALIZED(REF, DESTRUCT, __FILE__, __LINE__)
-
-static inline _PyStackRef
-PyStackRef_MakeHeapSafe(_PyStackRef ref)
+static inline void
+_PyStackRef_CLOSE_SPECIALIZED(_PyStackRef ref, destructor destruct, const char *filename, int linenumber)
 {
-    return ref;
+    assert(!PyStackRef_IsError(ref));
+    assert(!PyStackRef_IsNull(ref));
+    assert(!PyStackRef_IsTaggedInt(ref));
+    PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
+    if (PyStackRef_RefcountOnObject(ref)) {
+        _Py_DECREF_SPECIALIZED(obj, destruct);
+    }
 }
+#define PyStackRef_CLOSE_SPECIALIZED(REF, DESTRUCT) _PyStackRef_CLOSE_SPECIALIZED(REF, DESTRUCT, __FILE__, __LINE__)
 
 static inline _PyStackRef
-PyStackRef_Borrow(_PyStackRef ref)
+_PyStackRef_Borrow(_PyStackRef ref, const char *filename, int linenumber)
 {
-    return PyStackRef_DUP(ref);
+    assert(!PyStackRef_IsError(ref));
+    assert(!PyStackRef_IsNull(ref));
+    if (PyStackRef_IsTaggedInt(ref)) {
+        return ref;
+    }
+    PyObject *obj = _Py_stackref_get_object(ref);
+    return _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
 }
+#define PyStackRef_Borrow(REF) _PyStackRef_Borrow((REF), __FILE__, __LINE__)
 
 #define PyStackRef_CLEAR(REF) \
     do { \
@@ -234,27 +283,44 @@ static inline _PyStackRef
 _PyStackRef_FromPyObjectStealMortal(PyObject *obj, const char *filename, int linenumber)
 {
     assert(!_Py_IsImmortal(obj));
-    return _Py_stackref_create(obj, filename, linenumber);
+    return _Py_stackref_create(obj, 0, filename, linenumber);
 }
 #define PyStackRef_FromPyObjectStealMortal(obj) _PyStackRef_FromPyObjectStealMortal(_PyObject_CAST(obj), __FILE__, __LINE__)
 
 static inline bool
 PyStackRef_IsHeapSafe(_PyStackRef ref)
 {
-    return true;
+    if ((ref.index & Py_TAG_BITS) != Py_TAG_REFCNT || PyStackRef_IsNull(ref)) {
+        // Tagged ints and ERROR are included.
+        return true;
+    }
+
+    PyObject *obj = _Py_stackref_get_object(ref);
+    return _Py_IsImmortal(obj);
 }
 
+static inline _PyStackRef
+_PyStackRef_MakeHeapSafe(_PyStackRef ref, const char *filename, int linenumber)
+{
+    if (PyStackRef_IsHeapSafe(ref)) {
+        return ref;
+    }
+
+    PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
+    Py_INCREF(obj);
+    return _Py_stackref_create(obj, 0, filename, linenumber);
+}
+#define PyStackRef_MakeHeapSafe(REF) _PyStackRef_MakeHeapSafe(REF, __FILE__, __LINE__)
+
 static inline _PyStackRef
 _PyStackRef_FromPyObjectNewMortal(PyObject *obj, const char *filename, int linenumber)
 {
     assert(!_Py_IsStaticImmortal(obj));
-    Py_INCREF(obj);
-    return _Py_stackref_create(obj, filename, linenumber);
+    Py_XINCREF(obj);
+    return _Py_stackref_create(obj, 0, filename, linenumber);
 }
 #define PyStackRef_FromPyObjectNewMortal(obj) _PyStackRef_FromPyObjectNewMortal(_PyObject_CAST(obj), __FILE__, __LINE__)
 
-#define PyStackRef_RefcountOnObject(REF) 1
-
 extern int PyStackRef_Is(_PyStackRef a, _PyStackRef b);
 
 extern bool PyStackRef_IsTaggedInt(_PyStackRef ref);
@@ -287,6 +353,7 @@ PyStackRef_Wrap(void *ptr)
 {
     assert(ptr != NULL);
 #ifdef Py_DEBUG
+    assert(((uintptr_t)ptr & Py_TAG_BITS) == 0);
     return (_PyStackRef){ .bits = ((uintptr_t)ptr) | Py_TAG_INVALID };
 #else
     return (_PyStackRef){ .bits = (uintptr_t)ptr };

Misc/NEWS.d/next/Core_and_Builtins/2025-10-03-17-51-43.gh-issue-139475._684ED.rst

@@ -0,0 +1,2 @@
+Changes in stackref debugging mode when ``Py_STACKREF_DEBUG`` is set. We use
+the same pattern of refcounting for stackrefs as in production build.

Python/stackrefs.c

@@ -109,18 +109,19 @@ _Py_stackref_close(_PyStackRef ref, const char *filename, int linenumber)
 }
 
 _PyStackRef
-_Py_stackref_create(PyObject *obj, const char *filename, int linenumber)
+_Py_stackref_create(PyObject *obj, uint16_t flags, const char *filename, int linenumber)
 {
     if (obj == NULL) {
         Py_FatalError("Cannot create a stackref for NULL");
     }
     PyInterpreterState *interp = PyInterpreterState_Get();
     uint64_t new_id = interp->next_stackref;
-    interp->next_stackref = new_id + 2;
+    interp->next_stackref = new_id + 4;
     TableEntry *entry = make_table_entry(obj, filename, linenumber);
     if (entry == NULL) {
         Py_FatalError("No memory left for stackref debug table");
     }
+    new_id |= flags;
     if (_Py_hashtable_set(interp->open_stackrefs_table, (void *)new_id, entry) < 0) {
         Py_FatalError("No memory left for stackref debug table");
     }
@@ -194,24 +195,18 @@ _Py_stackref_report_leaks(PyInterpreterState *interp)
     }
 }
 
-void
-_PyStackRef_CLOSE_SPECIALIZED(_PyStackRef ref, destructor destruct, const char *filename, int linenumber)
-{
-    PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
-    _Py_DECREF_SPECIALIZED(obj, destruct);
-}
-
 _PyStackRef PyStackRef_TagInt(intptr_t i)
 {
-    return (_PyStackRef){ .index = (i << 1) + 1 };
+    assert(Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, (i << 2), 2) == i);
+    return (_PyStackRef){ .index = (i << 2) | Py_INT_TAG };
 }
 
 intptr_t
 PyStackRef_UntagInt(_PyStackRef i)
 {
     assert(PyStackRef_IsTaggedInt(i));
     intptr_t val = (intptr_t)i.index;
-    return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, val, 1);
+    return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, val, 2);
 }
 
 bool
@@ -223,8 +218,9 @@ PyStackRef_IsNullOrInt(_PyStackRef ref)
 _PyStackRef
 PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
 {
-    assert(ref.index <= INT_MAX - 2); // No overflow
-    return (_PyStackRef){ .index = ref.index + 2 };
+    assert(PyStackRef_IsTaggedInt(ref));
+    assert((ref.index & (~Py_TAG_BITS)) != (INT_MAX & (~Py_TAG_BITS))); // Isn't about to overflow
+    return (_PyStackRef){ .index = ref.index + 4 };
 }