Partially address review

Co-Authored-By: Sam Gross <655866+colesbury@users.noreply.github.com>
python · Fidget-Spinner · Jul 7, 2024 · Jul 7, 2024 · Sep 15, 2024 · Sep 15, 2024
commit ba19770f28c210109c9b6113d53e728f0848a888
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -121,7 +121,7 @@ PyAPI_FUNC(int) _PyDict_SetItem_KnownHash_LockHeld(PyDictObject *mp, PyObject *k
 PyAPI_FUNC(int) _PyDict_GetItemRef_KnownHash_LockHeld(PyDictObject *op, PyObject *key, Py_hash_t hash, PyObject **result);
 extern int _PyDict_GetItemRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, PyObject **result);
 extern int _PyDict_GetItemRef_Unicode_LockHeld(PyDictObject *op, PyObject *key, PyObject **result);
-extern int _PyDict_GetItem_KnownHash_StackRef(PyDictObject *op, PyObject *key, Py_hash_t hash, _PyStackRef *result);
+extern int _PyDict_GetItemStackRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, _PyStackRef *result);
 extern int _PyObjectDict_SetItem(PyTypeObject *tp, PyObject *obj, PyObject **dictptr, PyObject *name, PyObject *value);
 extern int _PyDict_GetItemStackRef(PyObject *op, PyObject *key, _PyStackRef *result);
 

diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
@@ -623,6 +623,24 @@ _Py_TryXGetRef(PyObject **ptr)
     return NULL;
 }
 
+// This belongs here rather than pycore_stackref.h because including pycore_object.h
+// there causes a circular include.
+static inline _PyStackRef
+_Py_TryXGetStackRef(PyObject **ptr)
+{
+    PyObject *value = _Py_atomic_load_ptr(ptr);
+    if (value == NULL) {
+        return PyStackRef_NULL;
+    }
+    if (_Py_IsImmortal(value) || _PyObject_HasDeferredRefcount(value)) {
+        return (_PyStackRef){ .bits = (uintptr_t)value | Py_TAG_DEFERRED };
+    }
+    if (_Py_TryIncrefCompare(ptr, value)) {
+        return PyStackRef_FromPyObjectSteal(value);
+    }
+    return PyStackRef_NULL;
+}
+
 /* Like Py_NewRef but also optimistically sets _Py_REF_MAYBE_WEAKREF
    on objects owned by a different thread. */
 static inline PyObject *

@@ -2337,7 +2337,7 @@ _PyDict_GetItemRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, Py
 }
 
 int
-_PyDict_GetItem_KnownHash_StackRef(PyDictObject *op, PyObject *key, Py_hash_t hash, _PyStackRef *result)
+_PyDict_GetItemStackRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash_t hash, _PyStackRef *result)
 {
     Py_ssize_t ix = _Py_dict_lookup_threadsafe_stackref(op, key, hash, result);
     assert(ix >= 0 || PyStackRef_IsNull(*result));
@@ -2384,7 +2384,7 @@ _PyDict_GetItemStackRef(PyObject *op, PyObject *key, _PyStackRef *result)
         return -1;
     }
 
-    return _PyDict_GetItem_KnownHash_StackRef((PyDictObject *)op, key, hash, result);
+    return _PyDict_GetItemStackRef_KnownHash((PyDictObject *)op, key, hash, result);
 }
 
 int
@@ -7148,14 +7148,8 @@ _PyObject_TryGetInstanceAttributeStackRef(PyObject *obj, PyObject *name, _PyStac
     }
 
 #ifdef Py_GIL_DISABLED
-    PyObject *value = _Py_atomic_load_ptr_acquire(&values->values[ix]);
-    if (value == NULL) {
-        *attr = PyStackRef_NULL;
-    }
-    else {
-        *attr = PyStackRef_FromPyObjectNew(value);
-    }
-    if (value == _Py_atomic_load_ptr_acquire(&values->values[ix])) {
+    *attr = _Py_TryXGetStackRef(&values->values[ix]);
+    if (PyStackRef_AsPyObjectBorrow(*attr) == _Py_atomic_load_ptr_acquire(&values->values[ix])) {
         return true;
     }
 
@@ -7170,13 +7164,7 @@ _PyObject_TryGetInstanceAttributeStackRef(PyObject *obj, PyObject *name, _PyStac
         if (dict == NULL) {
             // Still no dict, we can read from the values
             assert(values->valid);
-            value = values->values[ix];
-            if (value != NULL) {
-                *attr = PyStackRef_FromPyObjectNew(value);
-            }
-            else {
-                *attr = PyStackRef_NULL;
-            }
+            *attr = _Py_TryXGetStackRef(&values->values[ix]);
             success = true;
         }
 
@@ -7195,8 +7183,7 @@ _PyObject_TryGetInstanceAttributeStackRef(PyObject *obj, PyObject *name, _PyStac
     Py_BEGIN_CRITICAL_SECTION(dict);
 
     if (dict->ma_values == values && FT_ATOMIC_LOAD_UINT8(values->valid)) {
-        value = _Py_atomic_load_ptr_relaxed(&values->values[ix]);
-        *attr = PyStackRef_FromPyObjectNew(value);
+        *attr = _Py_TryXGetStackRef(&values->values[ix]);
         success = true;
     } else {
         // Caller needs to lookup from the dictionary

@@ -1636,7 +1636,13 @@ _PyObject_GetMethodStackRef(PyObject *obj, PyObject *name, _PyStackRef *method)
     }
 
     if (tp->tp_getattro != PyObject_GenericGetAttr || !PyUnicode_CheckExact(name)) {
-        *method = PyStackRef_FromPyObjectSteal(PyObject_GetAttr(obj, name));
+        PyObject *attr_o = PyObject_GetAttr(obj, name);
+        if (attr_o != NULL) {
+            *method = PyStackRef_FromPyObjectSteal(attr_o);
+        }
+        else {
+            *method = PyStackRef_NULL;
+        }
         return 0;
     }
 

@@ -5489,11 +5489,12 @@ find_name_in_mro_stackref(PyTypeObject *type, PyObject *name, int *error, _PySta
         PyObject *base = PyTuple_GET_ITEM(mro, i);
         PyObject *dict = lookup_tp_dict(_PyType_CAST(base));
         assert(dict && PyDict_Check(dict));
-        if (_PyDict_GetItem_KnownHash_StackRef((PyDictObject *)dict, name, hash, result) < 0) {
+        int code = _PyDict_GetItemStackRef_KnownHash((PyDictObject *)dict, name, hash, result);
+        if (code < 0) {
             *error = -1;
             goto done;
         }
-        if (!PyStackRef_IsNull(*result)) {
+        if (code == 1) {
             break;
         }
     }