moved free-threaded implementation details into a separate file

Author: wjakob

cmake/nanobind-config.cmake

@@ -175,6 +175,8 @@ function (nanobind_build_library TARGET_NAME)
     ${NB_DIR}/src/nb_enum.cpp
     ${NB_DIR}/src/nb_ndarray.cpp
     ${NB_DIR}/src/nb_static_property.cpp
+    ${NB_DIR}/src/nb_ft.h
+    ${NB_DIR}/src/nb_ft.cpp
     ${NB_DIR}/src/common.cpp
     ${NB_DIR}/src/error.cpp
     ${NB_DIR}/src/trampoline.cpp

src/common.cpp

@@ -1178,22 +1178,6 @@ bool issubclass(PyObject *a, PyObject *b) {
 
 // ========================================================================
 
-/// Make an object immortal when targeting free-threaded Python
-void maybe_make_immortal(PyObject *op) {
-#ifdef NB_FREE_THREADED
-    // See CPython's Objects/object.c
-    if (PyObject_IS_GC(op))
-        PyObject_GC_UnTrack(op);
-    op->ob_tid = _Py_UNOWNED_TID;
-    op->ob_ref_local = _Py_IMMORTAL_REFCNT_LOCAL;
-    op->ob_ref_shared = 0;
-#else
-    (void) op;
-#endif
-}
-
-// ========================================================================
-
 PyObject *dict_get_item_ref_or_fail(PyObject *d, PyObject *k) {
     PyObject *value;
     bool error = false;

src/nb_combined.cpp

@@ -78,6 +78,7 @@
 #include "nb_enum.cpp"
 #include "nb_ndarray.cpp"
 #include "nb_static_property.cpp"
+#include "nb_ft.cpp"
 #include "error.cpp"
 #include "common.cpp"
 #include "implicit.cpp"

src/nb_enum.cpp

@@ -1,4 +1,5 @@
 #include "nb_internals.h"
+#include "nb_ft.h"
 
 NAMESPACE_BEGIN(NB_NAMESPACE)
 NAMESPACE_BEGIN(detail)
@@ -95,7 +96,7 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
         #endif
     }
 
-    maybe_make_immortal(result.ptr());
+    make_immortal(result.ptr());
 
     result.attr("__nb_enum__") = capsule(t, [](void *p) noexcept {
         type_init_data *t = (type_init_data *) p;

src/nb_ft.cpp

@@ -0,0 +1,56 @@
+#include <nanobind/nanobind.h>
+#include "nb_ft.h"
+
+#if defined(Py_GIL_DISABLED)
+/// Make an object immortal when targeting free-threaded Python
+void make_immortal(PyObject *op) noexcept {
+    // See CPython's Objects/object.c
+    if (PyObject_IS_GC(op))
+        PyObject_GC_UnTrack(op);
+    op->ob_tid = _Py_UNOWNED_TID;
+    op->ob_ref_local = _Py_IMMORTAL_REFCNT_LOCAL;
+    op->ob_ref_shared = 0;
+}
+
+#if !(0 && PY_VERSION_HEX >= 0x030E00A5)
+void nb_enable_try_inc_ref(PyObject *obj) noexcept {
+    // Since this is called during object construction, we know that we have
+    // the only reference to the object and can use a non-atomic write.
+    assert(obj->ob_ref_shared == 0);
+    obj->ob_ref_shared = _Py_REF_MAYBE_WEAKREF;
+}
+
+bool nb_try_inc_ref(PyObject *obj) noexcept {
+    // See https://github.com/python/cpython/blob/d05140f9f77d7dfc753dd1e5ac3a5962aaa03eff/Include/internal/pycore_object.h#L761
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&obj->ob_ref_local);
+    local += 1;
+    if (local == 0) {
+        // immortal
+        return true;
+    }
+    if (_Py_IsOwnedByCurrentThread(obj)) {
+        _Py_atomic_store_uint32_relaxed(&obj->ob_ref_local, local);
+#ifdef Py_REF_DEBUG
+        _Py_INCREF_IncRefTotal();
+#endif
+        return true;
+    }
+    Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared);
+    for (;;) {
+        // If the shared refcount is zero and the object is either merged
+        // or may not have weak references, then we cannot incref it.
+        if (shared == 0 || shared == _Py_REF_MERGED) {
+            return false;
+        }
+
+        if (_Py_atomic_compare_exchange_ssize(
+                &obj->ob_ref_shared, &shared, shared + (1 << _Py_REF_SHARED_SHIFT))) {
+#ifdef Py_REF_DEBUG
+            _Py_INCREF_IncRefTotal();
+#endif
+            return true;
+        }
+    }
+}
+#endif
+#endif

src/nb_ft.h

@@ -0,0 +1,37 @@
+/*
+    src/nb_ft.h: implementation details related to free-threaded Python
+
+    Copyright (c) 2022 Wenzel Jakob
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
+#if !defined(Py_GIL_DISABLED)
+/// Trivial implementations for non-free-threaded Python
+inline void make_immortal(PyObject *) noexcept { }
+inline void nb_enable_try_inc_ref(PyObject *) noexcept { }
+inline bool nb_try_inc_ref(PyObject *obj) noexcept {
+    if (Py_REFCNT(obj) > 0) {
+        Py_INCREF(obj);
+        return true;
+    }
+    return false;
+}
+#else
+extern void make_immortal(PyObject *op) noexcept;
+
+#if 0 && PY_VERSION_HEX >= 0x030E00A5
+/// Sufficiently recent CPython versions provide an API for the following operations
+inline void nb_enable_try_inc_ref(PyObject *obj) noexcept {
+    PyUnstable_EnableTryIncRef(obj);
+}
+inline bool nb_try_inc_ref(PyObject *obj) noexcept {
+    return PyUnstable_TryIncRef(obj);
+}
+#else
+/// Otherwise, nanabind ships with a low-level implementation
+extern void nb_enable_try_inc_ref(PyObject *) noexcept;
+extern bool nb_try_inc_ref(PyObject *obj) noexcept;
+#endif
+#endif

src/nb_func.cpp

@@ -9,6 +9,7 @@
 
 #include "nb_internals.h"
 #include "buffer.h"
+#include "nb_ft.h"
 
 /// Maximum number of arguments supported by 'nb_vectorcall_simple'
 #define NB_MAXARGS_SIMPLE 8
@@ -289,7 +290,7 @@ PyObject *nb_func_new(const void *in_) noexcept {
     check(func, "nb::detail::nb_func_new(\"%s\"): alloc. failed (1).",
           name_cstr);
 
-    maybe_make_immortal((PyObject *) func);
+    make_immortal((PyObject *) func);
 
     // Check if the complex dispatch loop is needed
     bool complex_call = can_mutate_args || has_var_kwargs || has_var_args ||

src/nb_internals.h

@@ -513,7 +513,6 @@ struct lock_obj { lock_obj(PyObject *) { } };
 
 extern char *strdup_check(const char *);
 extern void *malloc_check(size_t size);
-extern void maybe_make_immortal(PyObject *op);
 
 extern char *extract_name(const char *cmd, const char *prefix, const char *s);
 

src/nb_type.cpp

@@ -8,6 +8,7 @@
 */
 
 #include "nb_internals.h"
+#include "nb_ft.h"
 
 #if defined(_MSC_VER)
 #  pragma warning(disable: 4706) // assignment within conditional expression
@@ -40,62 +41,6 @@ static PyObject **nb_weaklist_ptr(PyObject *self) {
     return weaklistoffset ? (PyObject **) ((uint8_t *) self + weaklistoffset) : nullptr;
 }
 
-static void nb_enable_try_inc_ref(PyObject *obj) noexcept {
-#if 0 && defined(Py_GIL_DISABLED) && PY_VERSION_HEX >= 0x030E00A5
-    PyUnstable_EnableTryIncRef(obj);
-#elif defined(Py_GIL_DISABLED)
-    // Since this is called during object construction, we know that we have
-    // the only reference to the object and can use a non-atomic write.
-    assert(obj->ob_ref_shared == 0);
-    obj->ob_ref_shared = _Py_REF_MAYBE_WEAKREF;
-#else
-    (void) obj;
-#endif
-}
-
-static bool nb_try_inc_ref(PyObject *obj) noexcept {
-#if 0 && defined(Py_GIL_DISABLED) && PY_VERSION_HEX >= 0x030E00A5
-    return PyUnstable_TryIncRef(obj);
-#elif defined(Py_GIL_DISABLED)
-    // See https://github.com/python/cpython/blob/d05140f9f77d7dfc753dd1e5ac3a5962aaa03eff/Include/internal/pycore_object.h#L761
-    uint32_t local = _Py_atomic_load_uint32_relaxed(&obj->ob_ref_local);
-    local += 1;
-    if (local == 0) {
-        // immortal
-        return true;
-    }
-    if (_Py_IsOwnedByCurrentThread(obj)) {
-        _Py_atomic_store_uint32_relaxed(&obj->ob_ref_local, local);
-#ifdef Py_REF_DEBUG
-        _Py_INCREF_IncRefTotal();
-#endif
-        return true;
-    }
-    Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared);
-    for (;;) {
-        // If the shared refcount is zero and the object is either merged
-        // or may not have weak references, then we cannot incref it.
-        if (shared == 0 || shared == _Py_REF_MERGED) {
-            return false;
-        }
-
-        if (_Py_atomic_compare_exchange_ssize(
-                &obj->ob_ref_shared, &shared, shared + (1 << _Py_REF_SHARED_SHIFT))) {
-#ifdef Py_REF_DEBUG
-            _Py_INCREF_IncRefTotal();
-#endif
-            return true;
-        }
-    }
-#else
-    if (Py_REFCNT(obj) > 0) {
-        Py_INCREF(obj);
-        return true;
-    }
-    return false;
-#endif
-}
-
 static PyGetSetDef inst_getset[] = {
     { "__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict, nullptr, nullptr },
     { nullptr, nullptr, nullptr, nullptr, nullptr }
@@ -945,7 +890,7 @@ static PyTypeObject *nb_type_tp(size_t supplement) noexcept {
         tp = (PyTypeObject *) nb_type_from_metaclass(
             internals_->nb_meta, internals_->nb_module, &spec);
 
-        maybe_make_immortal((PyObject *) tp);
+        make_immortal((PyObject *) tp);
 
         handle(tp).attr("__module__") = "nanobind";
 
@@ -1369,7 +1314,7 @@ PyObject *nb_type_new(const type_init_data *t) noexcept {
 
     Py_DECREF(metaclass);
 
-    maybe_make_immortal(result);
+    make_immortal(result);
 
     type_data *to = nb_type_data((PyTypeObject *) result);