Intrusive pointer support

This commit adds support for object hierarchies with builtin reference
counting. This provides an alternative to STL's `std::unique_ptr<>` /
`std::shared_ptr<>` that is more general *and* more efficient in binding
code.

.gitignore

@@ -1,7 +1,10 @@
 libnanobind-static.a
 libnanobind.dylib
 libnanobind.so
+libnanobind-abi3.dylib
+libnanobind-abi3.so
 nanobind.dll
+nanobind-abi3.dll
 
 /.ninja_deps
 /.ninja_log

README.md

@@ -264,8 +264,12 @@ changes are detailed below.
 - **Shared pointers and holders**. _nanobind_ removes the concept of a _holder
   type_, which caused inefficiencies and introduced complexity in _pybind11_.
   This has implications on object ownership, shared ownership, and interactions
-  with C++ shared/unique pointers. Please see the following [separate
-  document](docs/ownership.md) for the nitty-gritty details.
+  with C++ shared/unique pointers.
+
+  Please see the following [separate page](docs/ownership.md) for the
+  nitty-gritty details on shared and unique pointers. Classes with _intrusive_
+  reference counting also continue to be supported, please see the [linked
+  page](docs/intrusive.md) for details.
 
   The gist is that use of shared/unique pointers requires one or both of the
   following optional header files:

docs/intrusive.md

@@ -0,0 +1,74 @@
+# Intrusive reference counting
+
+Like _pybind11_, _nanobind_ provides a way of binding classes with builtin
+("intrusive") reference counting. This is the most general and cheapest way of
+handling shared ownership between C++ and Python, but it requires that the base
+class of an object hierarchy is adapted according to the needs of _nanobind_.
+
+Ordinarily, a simple class with intrusive reference counting might look as
+follows:
+
+```cpp
+class Object {
+public:
+    void inc_ref() const { ++m_ref_count; }
+
+    void dec_ref() const {
+        if (--m_ref_count == 0)
+            delete this;
+    }
+
+private:
+    mutable std::atomic<size_t> m_ref_count { 0 };
+};
+```
+
+The advantage of this over standard approaches like `std::shared_ptr<T>` is
+that no separate control block must be allocated. Subtle technical band-aids
+like `std::enable_shared_from_this<T>` to avoid undefined behavior are also
+no longer necessary.
+
+However, one issue that tends to arise when a type like `Object` is wrapped
+using _nanobind_ is that there are now *two* separate reference counts
+referring to the same object: one in Python's `PyObject`, and one in `Object`.
+This can lead to a problematic reference cycle:
+
+- Python's `PyObject` needs to keep `Object` alive so that the instance can be
+  safely passed to C++ functions.
+
+- The C++ `Object` may in turn need to keep the `PyObject` alive. This is the
+  case when a subclass uses `NB_TRAMPOLINE` and `NB_OVERRIDE` features to route
+  C++ virtual function calls back to a Python implementation.
+
+The source of the problem is that there are *two* separate counters that try to
+reason about the reference count of *one* instance. The solution is to reduce
+this to just one counter: 
+
+- if an instance lives purely on the C++ side, the `m_ref_count` field is
+  used to reason about the number of references.
+
+- The first time that an instance is exposed to Python (by being created from
+  Python, or by being returned from a bound C++ function), lifetime management
+  is delegated to Python.
+
+The files
+[`tests/object.h`](https://github.com/wjakob/nanobind/blob/master/tests/object.h)
+and
+[`tests/object.cpp`](https://github.com/wjakob/nanobind/blob/master/tests/object.cpp)
+contain an example implementation of a suitable base class named `Object`. It
+contains an extra optimization to use a single field of type
+`std::atomic<uintptr_t> m_state;` (8 bytes) to store *either* a reference
+counter or a pointer to a `PyObject*`.
+
+The main change in _nanobind_-based bindings is that the base class must
+specify a `nb::intrusive_ptr` annotation to inform an instance that lifetime
+management has been taken over by Python. This annotation is automatically
+inherited by all subclasses.
+
+```cpp
+nb::class_<Object>(
+    m, "Object",
+    nb::intrusive_ptr<Object>(
+        [](Object *o, PyObject *po) { o->set_self_py(po); }));
+```
+

docs/ownership.md

@@ -14,6 +14,14 @@ complexity; they were therefore removed in _nanobind_. This has implications on
 object ownership, shared ownership, and interactions with C++ shared/unique
 pointers.
 
+- **Intrusive reference counting**: Like _pybind11_, _nanobind_ provides a way
+  of binding classes with builtin ("intrusive") reference counting. This is the
+  most general and cheapest way of handling shared ownership between C++ and
+  Python, but it requires that the base class of an object hierarchy is adapted
+  according to the needs of _nanobind_. Details on using intrusive reference
+  counting can be found
+  [here](https://github.com/wjakob/nanobind/blob/master/docs/intrusive.md).
+
 - **Shared pointers**: It is possible to bind functions that receive and return
   `std::shared_ptr<T>` by including the optional type caster
   [`nanobind/stl/shared_ptr.h`](https://github.com/wjakob/nanobind/blob/master/include/nanobind/stl/shared_ptr.h)
@@ -23,13 +31,12 @@ pointers.
   ownership must be shared between Python and C++. _nanobind_ does this by
   increasing the reference count of the `PyObject` and then creating a
   `std::shared_ptr<T>` with a new control block containing a custom deleter
-  that reduces the Python reference count upon destruction of the shared
-  pointer.
+  that will in turn reduce the Python reference count upon destruction of the
+  shared pointer.
 
   When a C++ function returns a `std::shared_ptr<T>`, _nanobind_ checks if the
   instance already has a `PyObject` counterpart (nothing needs to be done in
-  this case). Otherwise, it creates a compact `PyObject` wrapping a pointer to
-  the instance data. It indicates shared ownership by creating a temporary
+  this case). Otherwise, it indicates shared ownership by creating a temporary
   `std::shared_ptr<T>` on the heap that will be destructed when the `PyObject`
   is garbage collected.
 

include/nanobind/nb_attr.h

@@ -57,6 +57,12 @@ struct is_enum { bool is_signed; };
 
 template <size_t /* Nurse */, size_t /* Patient */> struct keep_alive {};
 template <typename T> struct supplement {};
+template <typename T> struct intrusive_ptr {
+    intrusive_ptr(void (*set_self_py)(T *, PyObject *))
+        : set_self_py(set_self_py) {}
+    void (*set_self_py)(T *, PyObject *);
+};
+
 struct type_callback {
     type_callback(void (*value)(PyType_Slot **) noexcept) : value(value) {}
     void (*value)(PyType_Slot **) noexcept;

include/nanobind/nb_class.h

@@ -66,7 +66,10 @@ enum class type_flags : uint32_t {
     has_supplement           = (1 << 18),
 
     /// Instances of this type support dynamic attribute assignment
-    has_dynamic_attr         = (1 << 19)
+    has_dynamic_attr         = (1 << 19),
+
+    /// The class uses an intrusive reference counting approach
+    intrusive_ptr            = (1 << 20)
 };
 
 struct type_data {
@@ -87,6 +90,7 @@ struct type_data {
     bool (**implicit_py)(PyTypeObject *, PyObject *, cleanup_list *) noexcept;
     void (*type_callback)(PyType_Slot **) noexcept;
     void *supplement;
+    void (*set_self_py)(void *, PyObject *);
 #if defined(Py_LIMITED_API)
     size_t dictoffset;
 #endif
@@ -107,6 +111,12 @@ NB_INLINE void type_extra_apply(type_data &t, type_callback c) {
     t.type_callback = c.value;
 }
 
+template <typename T>
+NB_INLINE void type_extra_apply(type_data &t, intrusive_ptr<T> ip) {
+    t.flags |= (uint32_t) type_flags::intrusive_ptr;
+    t.set_self_py = (void (*)(void *, PyObject *)) ip.set_self_py;
+}
+
 NB_INLINE void type_extra_apply(type_data &t, is_enum e) {
     if (e.is_signed)
         t.flags |= (uint32_t) type_flags::is_signed_enum;

include/nanobind/stl/shared_ptr.h

@@ -1,3 +1,12 @@
+/*
+    nanobind/stl/shared_ptr.h: Type caster for std::shared_ptr<T>
+
+    Copyright (c) 2022 Wenzel Jakob <wenzel.jakob@epfl.ch>
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
 #pragma once
 
 #include <nanobind/nanobind.h>

include/nanobind/stl/unique_ptr.h

@@ -1,3 +1,12 @@
+/*
+    nanobind/stl/unique_ptr.h: Type caster for std::unique_ptr<T>
+
+    Copyright (c) 2022 Wenzel Jakob <wenzel.jakob@epfl.ch>
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
 #pragma once
 
 #include <nanobind/nanobind.h>

src/nb_func.cpp

@@ -567,6 +567,10 @@ static PyObject *nb_func_vectorcall_complex(PyObject *self,
                     nb_inst *self_arg_nb = (nb_inst *) self_arg;
                     self_arg_nb->destruct = true;
                     self_arg_nb->ready = true;
+
+                    const type_data *t = nb_type_data(Py_TYPE(self_arg));
+                    if (t->flags & (uint32_t) type_flags::intrusive_ptr)
+                        t->set_self_py(inst_ptr(self_arg_nb), self_arg);
                 }
 
                 goto done;
@@ -690,6 +694,10 @@ static PyObject *nb_func_vectorcall_simple(PyObject *self,
                     nb_inst *self_arg_nb = (nb_inst *) self_arg;
                     self_arg_nb->destruct = true;
                     self_arg_nb->ready = true;
+
+                    const type_data *t = nb_type_data(Py_TYPE(self_arg));
+                    if (t->flags & (uint32_t) type_flags::intrusive_ptr)
+                        t->set_self_py(inst_ptr(self_arg_nb), self_arg);
                 }
 
                 goto done;

src/nb_internals.cpp

@@ -62,8 +62,14 @@
 #  define NB_BUILD_ABI ""
 #endif
 
+#if defined(Py_LIMITED_API)
+#  define NB_LIMITED_API "_limited"
+#else
+#  define NB_LIMITED_API ""
+#endif
+
 #define NB_INTERNALS_ID "__nb_internals_v" \
-    NB_TOSTRING(NB_INTERNALS_VERSION) NB_COMPILER_TYPE NB_STDLIB NB_BUILD_ABI NB_BUILD_TYPE "__"
+    NB_TOSTRING(NB_INTERNALS_VERSION) NB_COMPILER_TYPE NB_STDLIB NB_BUILD_ABI NB_BUILD_TYPE NB_LIMITED_API "__"
 
 NAMESPACE_BEGIN(NB_NAMESPACE)
 NAMESPACE_BEGIN(detail)

src/nb_type.cpp

@@ -307,7 +307,8 @@ PyObject *nb_type_new(const type_data *t) noexcept {
          has_base_py       = t->flags & (uint32_t) type_flags::has_base_py,
          has_type_callback = t->flags & (uint32_t) type_flags::has_type_callback,
          has_supplement    = t->flags & (uint32_t) type_flags::has_supplement,
-         has_dynamic_attr  = t->flags & (uint32_t) type_flags::has_dynamic_attr;
+         has_dynamic_attr  = t->flags & (uint32_t) type_flags::has_dynamic_attr,
+         intrusive_ptr     = t->flags & (uint32_t) type_flags::intrusive_ptr;
 
     nb_internals &internals = internals_get();
     str name(t->name), qualname = name;
@@ -351,9 +352,10 @@ PyObject *nb_type_new(const type_data *t) noexcept {
         base = (PyObject *) it->second->type_py;
     }
 
+    type_data *tb = nullptr;
     if (base) {
         // Check if the base type already has dynamic attributes
-        type_data *tb = nb_type_data((PyTypeObject *) base);
+        tb = nb_type_data((PyTypeObject *) base);
         if (tb->flags & (uint32_t) type_flags::has_dynamic_attr)
             has_dynamic_attr = true;
 
@@ -486,6 +488,12 @@ PyObject *nb_type_new(const type_data *t) noexcept {
     type_data *to = nb_type_data((PyTypeObject *) result);
     *to = *t;
 
+    if (!intrusive_ptr && tb &&
+        (tb->flags & (uint32_t) type_flags::intrusive_ptr)) {
+        to->flags |= (uint32_t) type_flags::intrusive_ptr;
+        to->set_self_py = tb->set_self_py;
+    }
+
     to->name = name_copy;
     to->type_py = (PyTypeObject *) result;
 
@@ -776,9 +784,13 @@ PyObject *nb_type_put(const std::type_info *cpp_type, void *value,
     if (rvp == rv_policy::reference_internal && (!cleanup || !cleanup->self()))
         return nullptr;
 
-    bool store_in_obj = rvp == rv_policy::copy || rvp == rv_policy::move;
-
     type_data *t = it2->second;
+    const bool intrusive = t->flags & (uint32_t) type_flags::intrusive_ptr;
+    if (intrusive)
+        rvp = rv_policy::take_ownership;
+
+    const bool store_in_obj = rvp == rv_policy::copy || rvp == rv_policy::move;
+
     nb_inst *inst =
         (nb_inst *) inst_new_impl(t->type_py, store_in_obj ? nullptr : value);
     if (!inst)
@@ -837,6 +849,9 @@ PyObject *nb_type_put(const std::type_info *cpp_type, void *value,
     if (rvp == rv_policy::reference_internal)
         keep_alive((PyObject *) inst, cleanup->self());
 
+    if (intrusive)
+        t->set_self_py(new_value, (PyObject *) inst);
+
     return (PyObject *) inst;
 }
 

tests/CMakeLists.txt

@@ -4,6 +4,7 @@ nanobind_add_module(test_holders_ext test_holders.cpp)
 nanobind_add_module(test_stl_ext test_stl.cpp)
 nanobind_add_module(test_enum_ext test_enum.cpp)
 nanobind_add_module(test_tensor_ext test_tensor.cpp)
+nanobind_add_module(test_intrusive_ext test_intrusive.cpp object.cpp object.h)
 
 set(TEST_FILES
   test_functions.py
@@ -12,6 +13,7 @@ set(TEST_FILES
   test_stl.py
   test_enum.py
   test_tensor.py
+  test_intrusive.py
 )
 
 if (NOT (CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_CURRENT_BINARY_DIR) OR MSVC)

tests/object.cpp

@@ -0,0 +1,74 @@
+#include "object.h"
+#include <stdexcept>
+
+static void (*object_inc_ref_py)(PyObject *) = nullptr;
+static void (*object_dec_ref_py)(PyObject *) = nullptr;
+
+void Object::inc_ref() const {
+    uintptr_t value = m_state.load(std::memory_order_relaxed);
+
+    while (true) {
+        if (value & 1) {
+            if (!m_state.compare_exchange_weak(value,
+                                               value + 2,
+                                               std::memory_order_relaxed,
+                                               std::memory_order_relaxed))
+                continue;
+        } else {
+            object_inc_ref_py((PyObject *) value);
+        }
+
+        break;
+    }
+}
+
+void Object::dec_ref() const {
+    uintptr_t value = m_state.load(std::memory_order_relaxed);
+
+    while (true) {
+        if (value & 1) {
+            if (value == 1) {
+                throw std::runtime_error("Object::dec_ref(): reference count underflow!");
+            } else if (value == 3) {
+                delete this;
+            } else {
+                if (!m_state.compare_exchange_weak(value,
+                                                   value - 2,
+                                                   std::memory_order_relaxed,
+                                                   std::memory_order_relaxed))
+                    continue;
+            }
+        } else {
+            object_dec_ref_py((PyObject *) value);
+        }
+        break;
+    }
+}
+
+void Object::set_self_py(PyObject *o) {
+    uintptr_t value = m_state.load(std::memory_order_relaxed);
+    if (value & 1) {
+        value >>= 1;
+        for (uintptr_t i = 0; i < value; ++i)
+            object_inc_ref_py(o);
+
+        m_state.store((uintptr_t) o);
+    } else {
+        throw std::runtime_error("Object::set_self_py(): a Python object was already present!");
+    }
+    value = m_state.load(std::memory_order_relaxed);
+}
+
+PyObject *Object::self_py() const {
+    uintptr_t value = m_state.load(std::memory_order_relaxed);
+    if (value & 1)
+        return nullptr;
+    else
+        return (PyObject *) value;
+}
+
+void object_init_py(void (*object_inc_ref_py_)(PyObject *),
+                    void (*object_dec_ref_py_)(PyObject *)) {
+    object_inc_ref_py = object_inc_ref_py_;
+    object_dec_ref_py = object_dec_ref_py_;
+}