_PyObject_TryGetInstanceAttributeStackRef

Author: Fidget-Spinner

Include/internal/pycore_dict.h

@@ -120,6 +120,7 @@ extern int _PyDict_GetItemRef_KnownHash(PyDictObject *op, PyObject *key, Py_hash
 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 _PyObjectDict_SetItem(PyTypeObject *tp, PyObject *obj, PyObject **dictptr, PyObject *name, PyObject *value);
+extern int _PyDict_GetItemStackRef(PyObject *op, PyObject *key, _PyStackRef *result);
 
 extern int _PyDict_Pop_KnownHash(
     PyDictObject *dict,

Include/internal/pycore_object.h

@@ -790,7 +790,8 @@ extern int _PyObject_StoreInstanceAttribute(PyObject *obj,
                                             PyObject *name, PyObject *value);
 extern bool _PyObject_TryGetInstanceAttribute(PyObject *obj, PyObject *name,
                                               PyObject **attr);
-
+extern bool _PyObject_TryGetInstanceAttributeStackRef(PyObject *obj, PyObject *name,
+                                              _PyStackRef *attr);
 #ifdef Py_GIL_DISABLED
 #  define MANAGED_DICT_OFFSET    (((Py_ssize_t)sizeof(PyObject *))*-1)
 #  define MANAGED_WEAKREF_OFFSET (((Py_ssize_t)sizeof(PyObject *))*-2)

Objects/dictobject.c

@@ -2363,6 +2363,24 @@ PyDict_GetItemRef(PyObject *op, PyObject *key, PyObject **result)
     return _PyDict_GetItemRef_KnownHash((PyDictObject *)op, key, hash, result);
 }
 
+int
+_PyDict_GetItemStackRef(PyObject *op, PyObject *key, _PyStackRef *result)
+{
+    if (!PyDict_Check(op)) {
+        PyErr_BadInternalCall();
+        *result = PyStackRef_NULL;
+        return -1;
+    }
+
+    Py_hash_t hash = _PyObject_HashFast(key);
+    if (hash == -1) {
+        *result = PyStackRef_NULL;
+        return -1;
+    }
+
+    return _PyDict_GetItem_KnownHash_StackRef((PyDictObject *)op, key, hash, result);
+}
+
 int
 _PyDict_GetItemRef_Unicode_LockHeld(PyDictObject *op, PyObject *key, PyObject **result)
 {
@@ -7075,6 +7093,79 @@ _PyObject_TryGetInstanceAttribute(PyObject *obj, PyObject *name, PyObject **attr
 #endif
 }
 
+bool
+_PyObject_TryGetInstanceAttributeStackRef(PyObject *obj, PyObject *name, _PyStackRef *attr)
+{
+    assert(PyUnicode_CheckExact(name));
+    PyDictValues *values = _PyObject_InlineValues(obj);
+    if (!FT_ATOMIC_LOAD_UINT8(values->valid)) {
+        return false;
+    }
+
+    PyDictKeysObject *keys = CACHED_KEYS(Py_TYPE(obj));
+    assert(keys != NULL);
+    Py_ssize_t ix = _PyDictKeys_StringLookup(keys, name);
+    if (ix == DKIX_EMPTY) {
+        *attr = PyStackRef_NULL;
+        return true;
+    }
+
+#ifdef Py_GIL_DISABLED
+    PyObject *value = _Py_atomic_load_ptr_acquire(&values->values[ix]);
+    *attr = PyStackRef_FromPyObjectNew(value);
+    if (value == _Py_atomic_load_ptr_acquire(&values->values[ix])) {
+        return true;
+    }
+
+    PyDictObject *dict = _PyObject_GetManagedDict(obj);
+    if (dict == NULL) {
+        // No dict, lock the object to prevent one from being
+        // materialized...
+        bool success = false;
+        Py_BEGIN_CRITICAL_SECTION(obj);
+
+        dict = _PyObject_GetManagedDict(obj);
+        if (dict == NULL) {
+            // Still no dict, we can read from the values
+            assert(values->valid);
+            value = values->values[ix];
+            *attr = PyStackRef_FromPyObjectNew(value);
+            success = true;
+        }
+
+        Py_END_CRITICAL_SECTION();
+
+        if (success) {
+            return true;
+        }
+    }
+
+    // We have a dictionary, we'll need to lock it to prevent
+    // the values from being resized.
+    assert(dict != NULL);
+
+    bool success;
+    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);
+        success = true;
+    } else {
+        // Caller needs to lookup from the dictionary
+        success = false;
+    }
+
+    Py_END_CRITICAL_SECTION();
+
+    return success;
+#else
+    PyObject *value = values->values[ix];
+    *attr = PyStackRef_FromPyObjectNew(value);
+    return true;
+#endif
+}
+
 int
 _PyObject_IsInstanceDictEmpty(PyObject *obj)
 {

Objects/object.c

@@ -1620,11 +1620,10 @@ _PyObject_GetMethodStackRef(PyObject *obj, PyObject *name, _PyStackRef *method)
             }
         }
     }
-    PyObject *dict, *attr;
+    PyObject *dict;
     if ((tp->tp_flags & Py_TPFLAGS_INLINE_VALUES) &&
-        _PyObject_TryGetInstanceAttribute(obj, name, &attr)) {
-        if (attr != NULL) {
-            *method = PyStackRef_FromPyObjectSteal(attr);
+        _PyObject_TryGetInstanceAttributeStackRef(obj, name, method)) {
+        if (!PyStackRef_IsNull(*method)) {
             PyStackRef_XCLOSE(descr_st);
             return 0;
         }
@@ -1644,9 +1643,7 @@ _PyObject_GetMethodStackRef(PyObject *obj, PyObject *name, _PyStackRef *method)
     }
     if (dict != NULL) {
         Py_INCREF(dict);
-        PyObject *item;
-        if (PyDict_GetItemRef(dict, name, &item) != 0) {
-            *method = PyStackRef_FromPyObjectSteal(item);
+        if (_PyDict_GetItemStackRef(dict, name, method) != 0) {
             // found or error
             Py_DECREF(dict);
             PyStackRef_CLOSE(descr_st);