MakeHeapSafe rework, _Py_stackref_set_borrowed_from added (NOP)

Author: efimov-mikhail

Include/internal/pycore_stackref.h

@@ -63,6 +63,7 @@ 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, uint16_t flags, const char *filename, int linenumber);
 PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber);
+PyAPI_FUNC(void) _Py_stackref_set_borrowed_from(_PyStackRef ref, _PyStackRef borrowed_from, const char *filename, int linenumber);
 extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);
 
 static const _PyStackRef PyStackRef_NULL = { .index = 0 };
@@ -259,6 +260,7 @@ _PyStackRef_CLOSE_SPECIALIZED(_PyStackRef ref, destructor destruct, const char *
     assert(!PyStackRef_IsNull(ref));
     assert(!PyStackRef_IsTaggedInt(ref));
     PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
+    assert(Py_REFCNT(obj) > 0);
     if (PyStackRef_RefcountOnObject(ref)) {
         _Py_DECREF_SPECIALIZED(obj, destruct);
     }
@@ -274,7 +276,11 @@ _PyStackRef_Borrow(_PyStackRef ref, const char *filename, int linenumber)
         return ref;
     }
     PyObject *obj = _Py_stackref_get_object(ref);
-    return _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
+    _PyStackRef new_ref = _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
+    if (!_Py_IsImmortal(obj)) {
+        _Py_stackref_set_borrowed_from(new_ref, ref, filename, linenumber);
+    }
+    return new_ref;
 }
 #define PyStackRef_Borrow(REF) _PyStackRef_Borrow((REF), __FILE__, __LINE__)
 
@@ -310,13 +316,22 @@ PyStackRef_IsHeapSafe(_PyStackRef ref)
 static inline _PyStackRef
 _PyStackRef_MakeHeapSafe(_PyStackRef ref, const char *filename, int linenumber)
 {
-    if (PyStackRef_IsHeapSafe(ref)) {
+    // Special references that can't be closed.
+    if (ref.index < INITIAL_STACKREF_INDEX) {
         return ref;
     }
 
+    bool heap_safe = PyStackRef_IsHeapSafe(ref);
     PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
-    Py_INCREF(obj);
-    return _Py_stackref_create(obj, 0, filename, linenumber);
+    uint16_t flags = 0;
+    if (heap_safe) {
+        // Close old ref and create a new one with the same flags.
+        // This is necessary for correct borrow checking.
+        flags = ref.index & Py_TAG_BITS;
+    } else {
+        Py_INCREF(obj);
+    }
+    return _Py_stackref_create(obj, flags, filename, linenumber);
 }
 #define PyStackRef_MakeHeapSafe(REF) _PyStackRef_MakeHeapSafe(REF, __FILE__, __LINE__)
 

Python/stackrefs.c

@@ -152,6 +152,11 @@ _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber
     entry->linenumber_borrow = linenumber;
 }
 
+void
+_Py_stackref_set_borrowed_from(_PyStackRef ref, _PyStackRef borrowed_from, const char *filename, int linenumber)
+{
+    assert(!PyStackRef_IsError(ref));
+}
 
 void
 _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref)