gh-121459: Deferred LOAD_GLOBAL

Include/internal/pycore_ceval.h

@@ -271,7 +271,7 @@ PyAPI_FUNC(PyObject **) _PyObjectArray_FromStackRefArray(_PyStackRef *input, Py_
 PyAPI_FUNC(void) _PyObjectArray_Free(PyObject **array, PyObject **scratch);
 
 PyAPI_FUNC(PyObject *) _PyEval_GetANext(PyObject *aiter);
-PyAPI_FUNC(PyObject *) _PyEval_LoadGlobal(PyObject *globals, PyObject *builtins, PyObject *name);
+PyAPI_FUNC(void) _PyEval_LoadGlobalStackRef(PyObject *globals, PyObject *builtins, PyObject *name, _PyStackRef *writeto);
 PyAPI_FUNC(PyObject *) _PyEval_GetAwaitable(PyObject *iterable, int oparg);
 PyAPI_FUNC(PyObject *) _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *name);
 

Include/internal/pycore_dict.h

@@ -10,6 +10,7 @@ extern "C" {
 
 #include "pycore_object.h"               // PyManagedDictPointer
 #include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_SSIZE_ACQUIRE
+#include "pycore_stackref.h"             // _PyStackRef
 
 // Unsafe flavor of PyDict_GetItemWithError(): no error checking
 extern PyObject* _PyDict_GetItemWithError(PyObject *dp, PyObject *key);
@@ -100,10 +101,11 @@ extern void _PyDictKeys_DecRef(PyDictKeysObject *keys);
  */
 extern Py_ssize_t _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
 extern Py_ssize_t _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
+extern Py_ssize_t _Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr);
 
 extern Py_ssize_t _PyDict_LookupIndex(PyDictObject *, PyObject *);
 extern Py_ssize_t _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyObject *key);
-PyAPI_FUNC(PyObject *)_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObject *);
+PyAPI_FUNC(void)_PyDict_LoadGlobalStackRef(PyDictObject *, PyDictObject *, PyObject *, _PyStackRef *);
 
 /* Consumes references to key and value */
 PyAPI_FUNC(int) _PyDict_SetItem_Take2(PyDictObject *op, PyObject *key, PyObject *value);

Objects/dictobject.c

@@ -1496,6 +1496,97 @@ _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyOb
     return ix;
 }
 
+Py_ssize_t
+_Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
+{
+    PyDictKeysObject *dk;
+    DictKeysKind kind;
+    Py_ssize_t ix;
+
+    ensure_shared_on_read(mp);
+
+    dk = _Py_atomic_load_ptr(&mp->ma_keys);
+    kind = dk->dk_kind;
+
+    if (kind != DICT_KEYS_GENERAL) {
+        if (PyUnicode_CheckExact(key)) {
+            ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash);
+        }
+        else {
+            ix = unicodekeys_lookup_generic_threadsafe(mp, dk, key, hash);
+        }
+        if (ix == DKIX_KEY_CHANGED) {
+            goto read_failed;
+        }
+
+        if (ix >= 0) {
+            if (kind == DICT_KEYS_SPLIT) {
+                PyDictValues *values = _Py_atomic_load_ptr(&mp->ma_values);
+                if (values == NULL)
+                    goto read_failed;
+
+                uint8_t capacity = _Py_atomic_load_uint8_relaxed(&values->capacity);
+                if (ix >= (Py_ssize_t)capacity)
+                    goto read_failed;
+
+                *value_addr = PyStackRef_FromPyObjectNew(values->values[ix]);
+                if (PyStackRef_IsNull(*value_addr)) {
+                    goto read_failed;
+                }
+
+                if (values != _Py_atomic_load_ptr(&mp->ma_values)) {
+                    goto read_failed;
+                }
+            }
+            else {
+                *value_addr = PyStackRef_FromPyObjectNew(DK_UNICODE_ENTRIES(dk)[ix].me_value);
+                if (PyStackRef_IsNull(*value_addr)) {
+                    goto read_failed;
+                }
+
+                if (dk != _Py_atomic_load_ptr(&mp->ma_keys)) {
+                    goto read_failed;
+                }
+            }
+        }
+        else {
+            *value_addr = PyStackRef_NULL;
+        }
+    }
+    else {
+        ix = dictkeys_generic_lookup_threadsafe(mp, dk, key, hash);
+        if (ix == DKIX_KEY_CHANGED) {
+            goto read_failed;
+        }
+        if (ix >= 0) {
+            *value_addr = PyStackRef_FromPyObjectNew(DK_ENTRIES(dk)[ix].me_value);
+            if (PyStackRef_IsNull(*value_addr)) {
+                goto read_failed;
+            }
+            if (dk != _Py_atomic_load_ptr(&mp->ma_keys)) {
+                goto read_failed;
+            }
+        }
+        else {
+            *value_addr = PyStackRef_NULL;
+        }
+    }
+
+    return ix;
+
+    PyObject *value;
+read_failed:
+    // In addition to the normal races of the dict being modified the _Py_TryXGetRef
+    // can all fail if they don't yet have a shared ref count.  That can happen here
+    // or in the *_lookup_* helper.  In that case we need to take the lock to avoid
+    // mutation and do a normal incref which will make them shared.
+    Py_BEGIN_CRITICAL_SECTION(mp);
+    ix = _Py_dict_lookup(mp, key, hash, &value);
+    *value_addr = value == NULL ? PyStackRef_NULL : PyStackRef_FromPyObjectNew(value);
+    Py_END_CRITICAL_SECTION();
+    return ix;
+}
+
 #else   // Py_GIL_DISABLED
 
 Py_ssize_t
@@ -1506,6 +1597,20 @@ _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyOb
     return ix;
 }
 
+Py_ssize_t
+_Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
+{
+    PyObject *val;
+    Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, &val);
+    if (val != NULL) {
+        *value_addr = PyStackRef_FromPyObjectNew(val);
+    }
+    else {
+        *value_addr = PyStackRef_NULL;
+    }
+    return ix;
+}
+
 #endif
 
 int
@@ -2391,31 +2496,31 @@ _PyDict_GetItemStringWithError(PyObject *v, const char *key)
  * key hash failed, key comparison failed, ...). Return NULL if the key doesn't
  * exist. Return the value if the key exists.
  *
- * Returns a new reference.
+ * Stores a new stack reference.
  */
-PyObject *
-_PyDict_LoadGlobal(PyDictObject *globals, PyDictObject *builtins, PyObject *key)
+void
+_PyDict_LoadGlobalStackRef(PyDictObject *globals, PyDictObject *builtins, PyObject *key, _PyStackRef *res)
 {
     Py_ssize_t ix;
     Py_hash_t hash;
-    PyObject *value;
 
     hash = _PyObject_HashFast(key);
     if (hash == -1) {
-        return NULL;
+        *res = PyStackRef_NULL;
     }
 
     /* namespace 1: globals */
-    ix = _Py_dict_lookup_threadsafe(globals, key, hash, &value);
-    if (ix == DKIX_ERROR)
-        return NULL;
-    if (ix != DKIX_EMPTY && value != NULL)
-        return value;
+    ix = _Py_dict_lookup_threadsafe_stackref(globals, key, hash, res);
+    if (ix == DKIX_ERROR) {
+        *res = PyStackRef_NULL;
+    }
+    if (ix != DKIX_EMPTY && !PyStackRef_IsNull(*res)) {
+        return;
+    }
 
     /* namespace 2: builtins */
-    ix = _Py_dict_lookup_threadsafe(builtins, key, hash, &value);
-    assert(ix >= 0 || value == NULL);
-    return value;
+    ix = _Py_dict_lookup_threadsafe_stackref(builtins, key, hash, res);
+    assert(ix >= 0 || PyStackRef_IsNull(*res));
 }
 
 /* Consumes references to key and value */

Python/bytecodes.c

@@ -1438,21 +1438,21 @@ dummy_func(
 
         inst(LOAD_FROM_DICT_OR_GLOBALS, (mod_or_class_dict -- v)) {
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg);
-            PyObject *v_o;
-            int err = PyMapping_GetOptionalItem(PyStackRef_AsPyObjectBorrow(mod_or_class_dict), name, &v_o);
-            if (err < 0) {
+            PyObject *v_o = NULL;
+            if (PyMapping_GetOptionalItem(PyStackRef_AsPyObjectBorrow(mod_or_class_dict), name, &v_o) < 0) {
                 ERROR_NO_POP();
             }
             if (v_o == NULL) {
                 if (PyDict_CheckExact(GLOBALS())
                     && PyDict_CheckExact(BUILTINS()))
                 {
-                    v_o = _PyDict_LoadGlobal((PyDictObject *)GLOBALS(),
+                    _PyDict_LoadGlobalStackRef((PyDictObject *)GLOBALS(),
                                             (PyDictObject *)BUILTINS(),
-                                            name);
-                    if (v_o == NULL) {
+                                            name,
+                                            STACK_ENTRY(v));
+                    if (PyStackRef_IsNull(v)) {
                         if (!_PyErr_Occurred(tstate)) {
-                            /* _PyDict_LoadGlobal() returns NULL without raising
+                            /* _PyDict_LoadGlobalStackRef() sets NULL without raising
                             * an exception if the key doesn't exist */
                             _PyEval_FormatExcCheckArg(tstate, PyExc_NameError,
                                                     NAME_ERROR_MSG, name);
@@ -1476,8 +1476,10 @@ dummy_func(
                     }
                 }
             }
+            if (v_o != NULL) {
+                v = PyStackRef_FromPyObjectSteal(v_o);
+            }
             DECREF_INPUTS();
-            v = PyStackRef_FromPyObjectSteal(v_o);
         }
 
         inst(LOAD_NAME, (-- v)) {
@@ -1507,10 +1509,9 @@ dummy_func(
 
         op(_LOAD_GLOBAL, ( -- res, null if (oparg & 1))) {
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg>>1);
-            PyObject *res_o = _PyEval_LoadGlobal(GLOBALS(), BUILTINS(), name);
-            ERROR_IF(res_o == NULL, error);
+            _PyEval_LoadGlobalStackRef(GLOBALS(), BUILTINS(), name, STACK_ENTRY(res));
+            ERROR_IF(PyStackRef_IsNull(res), error);
             null = PyStackRef_NULL;
-            res = PyStackRef_FromPyObjectSteal(res_o);
         }
 
         macro(LOAD_GLOBAL) =

Python/ceval.c

@@ -3072,15 +3072,14 @@ _PyEval_GetANext(PyObject *aiter)
     return awaitable;
 }
 
-PyObject *
-_PyEval_LoadGlobal(PyObject *globals, PyObject *builtins, PyObject *name)
+void
+_PyEval_LoadGlobalStackRef(PyObject *globals, PyObject *builtins, PyObject *name, _PyStackRef *writeto)
 {
-    PyObject *res;
     if (PyDict_CheckExact(globals) && PyDict_CheckExact(builtins)) {
-        res = _PyDict_LoadGlobal((PyDictObject *)globals,
+        _PyDict_LoadGlobalStackRef((PyDictObject *)globals,
                                     (PyDictObject *)builtins,
-                                    name);
-        if (res == NULL && !PyErr_Occurred()) {
+                                    name, writeto);
+        if (PyStackRef_IsNull(*writeto) && !PyErr_Occurred()) {
             /* _PyDict_LoadGlobal() returns NULL without raising
                 * an exception if the key doesn't exist */
             _PyEval_FormatExcCheckArg(PyThreadState_GET(), PyExc_NameError,
@@ -3090,22 +3089,25 @@ _PyEval_LoadGlobal(PyObject *globals, PyObject *builtins, PyObject *name)
     else {
         /* Slow-path if globals or builtins is not a dict */
         /* namespace 1: globals */
+        PyObject *res;
         if (PyMapping_GetOptionalItem(globals, name, &res) < 0) {
-            return NULL;
+            *writeto = PyStackRef_NULL;
+            return;
         }
         if (res == NULL) {
             /* namespace 2: builtins */
             if (PyMapping_GetOptionalItem(builtins, name, &res) < 0) {
-                return NULL;
+                *writeto = PyStackRef_NULL;
+                return;
             }
             if (res == NULL) {
                 _PyEval_FormatExcCheckArg(
                             PyThreadState_GET(), PyExc_NameError,
                             NAME_ERROR_MSG, name);
             }
         }
+        *writeto = PyStackRef_FromPyObjectSteal(res);
     }
-    return res;
 }
 
 PyObject *