Fix segfault in GC traversal of Python subclasses with weak references (#1206)

When ``nb::is_weak_referenceable()`` is used, nanobind installs ``tp_traverse``
and ``tp_clear`` callbacks to handle garbage collection of instance dictionaries
and weak reference lists.

When a Python subclass is created, Python may add its own instance dictionary
(potentially using managed dictionaries on Python 3.12+) or weak reference
list. Python's ``subtype_traverse`` function walks up the MRO and calls our
``tp_traverse`` callback. However, our callback was reading ``tp_dictoffset``
and ``tp_weaklistoffset`` directly from the type object, which includes
Python's additions. This caused crashes because:

1. On Python 3.11+, negative offsets are used (measured from the end of
   the object rather than the beginning) which we didn't handle
2. On Python 3.12+, managed dictionaries require special APIs that we weren't
   using and that aren't even available in the stable ABI/limited API

The fix is to always cache ``dictoffset``/``weaklistoffset`` in ``type_data``
when creating nanobind types, and use these cached values in ``nb_dict_ptr()``
and ``nb_weaklist_ptr()``. This ensures we only access dicts and weaklists
that nanobind created, while Python's ``subtype_traverse`` handles any
additions made by Python subclasses.

This requires an ABI bump due to the addition of ``dictoffset``/``weaklistoffset``
fields to ``type_data`` unconditionally (previously only in ``Py_LIMITED_API``
builds). Fortunately, the ABI was just bumped in commit aa1c9fd.

The commit also removes some redundant ``Py_TYPE()`` calls that have a
nontrivial cost on stable ABI builds.

Fixes #1201

Author: wjakob

include/nanobind/nb_class.h

@@ -128,10 +128,8 @@ struct type_data {
     };
     void (*set_self_py)(void *, PyObject *) noexcept;
     bool (*keep_shared_from_this_alive)(PyObject *) noexcept;
-#if defined(Py_LIMITED_API)
     uint32_t dictoffset;
     uint32_t weaklistoffset;
-#endif
 };
 
 /// Information about a type that is only relevant when it is being created

src/nb_type.cpp

@@ -21,23 +21,13 @@
 NAMESPACE_BEGIN(NB_NAMESPACE)
 NAMESPACE_BEGIN(detail)
 
-static PyObject **nb_dict_ptr(PyObject *self) {
-    PyTypeObject *tp = Py_TYPE(self);
-#if defined(Py_LIMITED_API)
+static PyObject **nb_dict_ptr(PyObject *self, PyTypeObject *tp) {
     Py_ssize_t dictoffset = nb_type_data(tp)->dictoffset;
-#else
-    Py_ssize_t dictoffset = tp->tp_dictoffset;
-#endif
     return dictoffset ? (PyObject **) ((uint8_t *) self + dictoffset) : nullptr;
 }
 
-static PyObject **nb_weaklist_ptr(PyObject *self) {
-    PyTypeObject *tp = Py_TYPE(self);
-#if defined(Py_LIMITED_API)
+static PyObject **nb_weaklist_ptr(PyObject *self, PyTypeObject *tp) {
     Py_ssize_t weaklistoffset = nb_type_data(tp)->weaklistoffset;
-#else
-    Py_ssize_t weaklistoffset = tp->tp_weaklistoffset;
-#endif
     return weaklistoffset ? (PyObject **) ((uint8_t *) self + weaklistoffset) : nullptr;
 }
 
@@ -47,18 +37,20 @@ static PyGetSetDef inst_getset[] = {
 };
 
 static int inst_clear(PyObject *self) {
-    PyObject **dict = nb_dict_ptr(self);
+    PyTypeObject *tp = Py_TYPE(self);
+    PyObject **dict = nb_dict_ptr(self, tp);
     if (dict)
         Py_CLEAR(*dict);
     return 0;
 }
 
 static int inst_traverse(PyObject *self, visitproc visit, void *arg) {
-    PyObject **dict = nb_dict_ptr(self);
+    PyTypeObject *tp = Py_TYPE(self);
+    PyObject **dict = nb_dict_ptr(self, tp);
     if (dict)
         Py_VISIT(*dict);
 #if PY_VERSION_HEX >= 0x03090000
-    Py_VISIT(Py_TYPE(self));
+    Py_VISIT(tp);
 #endif
     return 0;
 }
@@ -227,16 +219,16 @@ static void inst_dealloc(PyObject *self) {
         PyObject_GC_UnTrack(self);
 
         if (t->flags & (uint32_t) type_flags::has_dynamic_attr) {
-            PyObject **dict = nb_dict_ptr(self);
+            PyObject **dict = nb_dict_ptr(self, tp);
             if (dict)
                 Py_CLEAR(*dict);
         }
     }
 
     if (t->flags & (uint32_t) type_flags::is_weak_referenceable &&
-        nb_weaklist_ptr(self) != nullptr) {
+        nb_weaklist_ptr(self, tp) != nullptr) {
 #if defined(PYPY_VERSION)
-        PyObject **weaklist = nb_weaklist_ptr(self);
+        PyObject **weaklist = nb_weaklist_ptr(self, tp);
         if (weaklist)
             Py_CLEAR(*weaklist);
 #else
@@ -1362,13 +1354,10 @@ PyObject *nb_type_new(const type_init_data *t) noexcept {
     if (is_weak_referenceable)
         to->flags |= (uint32_t) type_flags::is_weak_referenceable;
 
-    #if defined(Py_LIMITED_API)
-        /* These must be set unconditionally so that nb_dict_ptr() /
-           nb_weaklist_ptr() return null (rather than garbage) on
-           objects whose types don't use the corresponding feature. */
-        to->dictoffset = (uint32_t) dictoffset;
-        to->weaklistoffset = (uint32_t) weaklistoffset;
-    #endif
+    // Always cache dictoffset/weaklistoffset so nb_dict_ptr()/nb_weaklist_ptr()
+    // only access dicts/weaklists created by nanobind, not those added by Python
+    to->dictoffset = (uint32_t) dictoffset;
+    to->weaklistoffset = (uint32_t) weaklistoffset;
 
     if (t->scope != nullptr)
         setattr(t->scope, t_name, result);

tests/test_classes.cpp

@@ -644,6 +644,11 @@ NB_MODULE(test_classes_ext, m) {
                nb::is_weak_referenceable(), nb::dynamic_attr())
         .def(nb::init<int>());
 
+    // test50_weakref_with_slots_subclass
+    struct StructWithWeakrefsOnly : Struct { };
+    nb::class_<StructWithWeakrefsOnly, Struct>(m, "StructWithWeakrefsOnly", nb::is_weak_referenceable())
+        .def(nb::init<int>());
+
     union Union {
         int i;
         float f;

tests/test_classes.py

@@ -945,3 +945,38 @@ def my_init(self):
 def test49_static_property_override():
     assert t.StaticPropertyOverride.x == 42
     assert t.StaticPropertyOverride2.x == 43
+
+def test50_weakref_with_slots_subclass():
+    """
+    Test that Python subclasses work correctly with nb::is_weak_referenceable()
+    base classes. The nb::is_weak_referenceable() flag causes nanobind to
+    install tp_traverse/tp_clear callbacks. When Python subclasses add their
+    own instance dictionaries (e.g., via managed dicts on Python 3.12+),
+    subtype_traverse calls our tp_traverse. We must only traverse dicts/weaklists
+    created by nanobind, not those added by Python.
+
+    Regression test for issue #1201.
+    """
+    import gc
+
+    # Create a Python subclass with __slots__
+    class SubClass(t.StructWithWeakrefsOnly):
+        __slots__ = 'hello',
+
+    # Create a sub-subclass without __slots__ (which should get a __dict__)
+    class SubSubClass(SubClass):
+        pass
+
+    # This should not crash
+    x = SubSubClass(42)
+    x.bye = 'blah'
+    assert x.value() == 42
+    assert x.bye == 'blah'
+
+    # Trigger GC to ensure inst_traverse doesn't crash
+    gc.collect()
+    gc.collect()
+
+    # Clean up
+    del x
+    gc.collect()

tests/test_classes_ext.pyi.ref

@@ -281,6 +281,9 @@ class StructWithWeakrefs(Struct):
 class StructWithWeakrefsAndDynamicAttrs(Struct):
     def __init__(self, arg: int, /) -> None: ...
 
+class StructWithWeakrefsOnly(Struct):
+    def __init__(self, arg: int, /) -> None: ...
+
 class Union:
     def __init__(self) -> None: ...