Repository: numpy/numpy
File: numpy/_core/src/umath/dispatching.cpp (C++)

Code before:
        Py_XDECREF(new_op_dtypes[i]);
    }

    /* Add this to the cache using the original types: */
    if (cacheable && PyArrayIdentityHash_SetItem(
                (PyArrayIdentityHash *)ufunc->_dispatch_cache,
                (PyObject **)op_dtypes, info, 0) < 0) {
        return NULL;
    }
    return info;
}

#ifdef Py_GIL_DISABLED
/*
 * Fast path for promote_and_get_info_and_ufuncimpl.
 * Acquires a read lock to check for a cache hit and then
 * only acquires a write lock on a cache miss to fill the cache
 */
static inline PyObject *
try_promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
        PyArrayObject *const ops[],
        PyArray_DTypeMeta *signature[],
        PyArray_DTypeMeta *op_dtypes[],
        npy_bool legacy_promotion_is_possible)
{
    ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache)->mutex)->lock_shared();
    PyObject *info = PyArrayIdentityHash_GetItem(
            (PyArrayIdentityHash *)ufunc->_dispatch_cache,
            (PyObject **)op_dtypes);
    ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache)->mutex)->unlock_shared();

    if (info != NULL && PyObject_TypeCheck(
                    PyTuple_GET_ITEM(info, 1), &PyArrayMethod_Type)) {
        /* Found the ArrayMethod and NOT a promoter: return it */
        return info;
    }

    // cache miss, need to acquire a write lock and recursively calculate the
    // correct dispatch resolution
    ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache)->mutex)->lock();
    info = promote_and_get_info_and_ufuncimpl(ufunc,
            ops, signature, op_dtypes, legacy_promotion_is_possible);
    ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache)->mutex)->unlock();

    return info;
}
#endif

/**
 * The central entry-point for the promotion and dispatching machinery.
 *

Review comment (suggestion):
it would be a little clearer to define a `std::shared_mutex *mutex` local variable and use that here and below so this complicated pointer casting only happens once

Code after:
        Py_XDECREF(new_op_dtypes[i]);
    }

    /* Add this to the cache using the original types: */
    if (cacheable && PyArrayIdentityHash_SetItem(
                (PyArrayIdentityHash *)ufunc->_dispatch_cache,
                (PyObject **)op_dtypes, info, 0) < 0) {
        return NULL;
    }
    return info;
}

#ifdef Py_GIL_DISABLED
/*
 * Fast path for promote_and_get_info_and_ufuncimpl.
 * Acquires a read lock to check for a cache hit and then
 * only acquires a write lock on a cache miss to fill the cache
 */
static inline PyObject *
promote_and_get_info_and_ufuncimpl_with_locking(
        PyUFuncObject *ufunc,
        PyArrayObject *const ops[],
        PyArray_DTypeMeta *signature[],
        PyArray_DTypeMeta *op_dtypes[],
        npy_bool legacy_promotion_is_possible)
{
    std::shared_mutex *mutex = ((std::shared_mutex *)((PyArrayIdentityHash *)ufunc->_dispatch_cache)->mutex);
    mutex->lock_shared();
    PyObject *info = PyArrayIdentityHash_GetItem(
            (PyArrayIdentityHash *)ufunc->_dispatch_cache,
            (PyObject **)op_dtypes);
    mutex->unlock_shared();

    if (info != NULL && PyObject_TypeCheck(
                    PyTuple_GET_ITEM(info, 1), &PyArrayMethod_Type)) {
        /* Found the ArrayMethod and NOT a promoter: return it */
        return info;
    }

    // cache miss, need to acquire a write lock and recursively calculate the
    // correct dispatch resolution
    mutex->lock();
    info = promote_and_get_info_and_ufuncimpl(ufunc,
            ops, signature, op_dtypes, legacy_promotion_is_possible);
    mutex->unlock();

    return info;
}
#endif

/**