Support enable_shared_from_this (#212)

This commit removes existing limitations related to ``std::shared_ptr`` for classes that inherit from ``std::enable_shared_from_this``.

docs/changelog.rst

@@ -34,7 +34,7 @@ Version 1.3.0 (TBD)
 * Reduced the per-instance overhead of nanobind by 1 pointer and simplified the
   internal hash table types to crunch ``libnanobind``. (commit `de018d
   <https://github.com/wjakob/nanobind/commit/de018db2d17905564703f1ade4aa201a22f8551f>`__).
-* Reduced the size of nanobind type objects by 6 pointers. (PR `#194
+* Reduced the size of nanobind type objects by 5 pointers. (PR `#194
   <https://github.com/wjakob/nanobind/pull/194>`__, `#195
   <https://github.com/wjakob/nanobind/pull/195>`__, and commit `d82ca9
   <https://github.com/wjakob/nanobind/commit/d82ca9c14191e74dd35dd5bf15fc90f5230319fb>`__).
@@ -108,6 +108,11 @@ Version 1.3.0 (TBD)
     ``some_enum < None`` will still fail, but now with a more
     informative error.
 
+* nanobind now has limited support for binding types that inherit from
+  ``std::enable_shared_from_this<T>``. See the :ref:`advanced section
+  on object ownership <enable_shared_from_this>` for more details.
+  (PR `#212 <https://github.com/wjakob/nanobind/pull/212>`__).
+
 * ABI version 8.
 
 Version 1.2.0 (April 24, 2023)

docs/ownership.rst

@@ -330,12 +330,30 @@ nanobind's support for shared pointers requires an extra include directive:
 You don't need to specify a return value policy annotation when a function
 returns a shared pointer.
 
-Shared pointer support has one major limitation in nanobind: the
-``std::enable_shared_from_this<T>`` base class that normally enables safe
-conversion of raw pointers to the associated shared pointer *may not be used*.
-Further detail can be found in the *advanced* :ref:`section <shared_ptr_adv>`
-on object ownership. If you need this feature, switch to intrusive reference
-counting explained below.
+nanobind's implementation of ``std::shared_ptr`` support typically
+allocates a new ``shared_ptr`` control block each time a Python object
+must be converted to ``std::shared_ptr<T>``. The new ``shared_ptr``
+"owns" a reference to the Python object, and its deleter drops that
+reference.  This has the advantage that the Python portion of the
+object will be kept alive by its C++-side references (which is
+important when implementing C++ virtual methods in Python), but it can
+be inefficient when passing the same object back and forth between
+Python and C++ many times, and it means that the ``use_count()``
+method of ``std::shared_ptr`` will return a value that does not
+capture all uses. Some of these problems can be mitigated by modifying
+``T`` so that it inherits from ``std::enable_shared_from_this<T>``.
+See the :ref:`advanced section <shared_ptr_adv>` on object ownership
+for more details on the implementation.
+
+nanobind has limited support for objects that inherit from
+``std::enable_shared_from_this<T>`` to allow safe conversion of raw
+pointers to shared pointers. The safest way to deal with these objects
+is to always use ``std::make_shared<T>(...)`` when constructing them in C++,
+and always pass them across the Python/C++ boundary wrapped in an explicit
+``std::shared_ptr<T>``. If you do this, then there shouldn't be any
+surprises. If you will be passing raw ``T*`` pointers around, then
+read the :ref:`advanced section on object ownership <enable_shared_from_this>`
+for additional caveats.
 
 .. _intrusive_intro:
 

docs/ownership_adv.rst

@@ -118,7 +118,7 @@ in the introductory section on object ownership and provides detail on how
 shared pointer conversion is *implemented* by nanobind.
 
 When the user calls a C++ function taking an argument of type
-``std::shared<T>`` from Python, ownership of that object must be
+``std::shared_ptr<T>`` from Python, ownership of that object must be
 shared between C++ to Python. 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 will in turn reduce the Python
@@ -139,19 +139,110 @@ true global reference count.
 
 .. _enable_shared_from_this:
 
-Limitations
-^^^^^^^^^^^
-
-nanobind refuses conversion of classes that derive from
-``std::enable_shared_from_this<T>``. This is a fundamental limitation:
-nanobind instances do not create a base shared pointer that declares
-ownership of an object. Other parts of a C++ codebase might then incorrectly
-assume ownership and eventually try to ``delete`` a nanobind instance
-allocated using ``pymalloc`` (which is undefined behavior). A compile-time
-assertion catches this and warns about the problem.
-
-Replacing shared pointers with :ref:`intrusive reference counting
-<intrusive>` fixes this limitations.
+enable_shared_from_this
+^^^^^^^^^^^^^^^^^^^^^^^
+
+The C++ standard library class ``std::enable_shared_from_this<T>``
+allows an object that inherits from it to locate an existing
+``std::shared_ptr<T>`` that manages that object. nanobind supports
+types that inherit from ``enable_shared_from_this``, with some caveats
+described in this section.
+
+Background (not nanobind-specific): Suppose a type ``ST`` inherits
+from ``std::enable_shared_from_this<ST>``. When a raw pointer ``ST
+*obj`` or ``std::unique_ptr<ST> obj`` is wrapped in a shared pointer
+using a constructor of the form ``std::shared_ptr<ST>(obj, ...)``, a
+reference to the new ``shared_ptr``\'s control block is saved (as
+``std::weak_ptr<ST>``) inside the object. This allows new
+``shared_ptr``\s that share ownership with the existing one to be
+obtained for the same object using ``obj->shared_from_this()`` or
+``obj->weak_from_this()``.
+
+nanobind's support for ``std::enable_shared_from_this`` consists of three
+behaviors:
+
+* If a raw pointer ``ST *obj`` is returned from C++ to Python, and
+  there already exists an associated ``std::shared_ptr<ST>`` which
+  ``obj->shared_from_this()`` can locate, then nanobind will produce a
+  Python instance that shares ownership with it. The behavior is
+  identical to what would happen if the C++ code did ``return
+  obj->shared_from_this();`` (returning an explicit
+  ``std::shared_ptr<ST>`` to Python) rather than ``return obj;``.
+  The return value policy has limited effect in this case; you will get
+  shared ownership on the Python side regardless of whether you used
+  `rv_policy::take_ownership` or `rv_policy::reference`.
+  (`rv_policy::copy` and `rv_policy::move` will still create a new
+  object that has no ongoing relationship to the returned pointer.)
+
+  * Note that this behavior occurs only if such a ``std::shared_ptr<ST>``
+    already exists! If not, then nanobind behaves as it would without
+    ``enable_shared_from_this``: a raw pointer will transfer exclusive
+    ownership to Python by default, or will create a non-owning reference
+    if you use `rv_policy::reference`.
+
+* If a Python object is passed to C++ as ``std::shared_ptr<ST> obj``,
+  and there already exists an associated ``std::shared_ptr<ST>`` which
+  ``obj->shared_from_this()`` can locate, then nanobind will produce a
+  ``std::shared_ptr<ST>`` that shares ownership with it: an additional
+  reference to the same control block, rather than a new control block
+  (as would occur without ``enable_shared_from_this``). This improves
+  performance and makes the result of ``shared_ptr::use_count()`` more
+  accurate.
+
+* If a Python object is passed to C++ as ``std::shared_ptr<ST> obj``, and
+  there is no associated ``std::shared_ptr<ST>`` that
+  ``obj->shared_from_this()`` can locate, then nanobind will produce
+  a ``std::shared_ptr<ST>`` as usual (with a new control block whose deleter
+  drops a Python object reference), *and* will do so in a way that enables
+  future calls to ``obj->shared_from_this()`` to find it as long
+  as any ``shared_ptr`` that shares this control block is still alive on
+  the C++ side.
+
+  (Once all of the ``std::shared_ptr<ST>``\s that share this control block
+  have been destroyed, the underlying PyObject reference being
+  managed by the ``shared_ptr`` deleter will be dropped,
+  and ``shared_from_this()`` will stop working. It can be reenabled by
+  passing the Python object back to C++ as ``std::shared_ptr<ST>`` once more,
+  which will create another control block.)
+
+Bindings for a class that supports ``enable_shared_from_this`` will be
+slightly larger than bindings for a class that doesn't, as nanobind
+must produce type-specific code to implement the above behaviors.
+
+.. warning:: The ``shared_from_this()`` method will only work when there
+   is actually a ``std::shared_ptr`` managing the object. A nanobind
+   instance constructed from Python will not have an associated
+   ``std::shared_ptr`` yet, so ``shared_from_this()`` will throw an
+   exception if you pass such an instance to C++ using a reference or
+   raw pointer. ``shared_from_this()`` will only work when there exists
+   a corresponding live ``std::shared_ptr`` on the C++ side.
+
+   The only situation where nanobind will create the first
+   ``std::shared_ptr`` for an object (thus enabling
+   ``shared_from_this()``), even with ``enable_shared_from_this``, is
+   when a Python instance is passed to C++ as the explicit type
+   ``std::shared_ptr<T>``. If you don't do this, or if no such
+   ``std::shared_ptr`` is still alive, then ``shared_from_this()`` will
+   throw an exception. It also works to create the ``std::shared_ptr``
+   on the C++ side, such as by using a factory function which always
+   uses ``std::make_shared<T>(...)`` to construct the object, and
+   returns the resulting ``std::shared_ptr<T>`` to Python.
+
+There is no way to enable ``shared_from_this`` immediately upon
+regular Python-side object construction (i.e., ``SomeType(*args)``
+rather than ``SomeType.some_fn(*args)``). If this limitation creates
+a problem for your application, you might get better results by using
+:ref:`intrusive reference counting <intrusive>` instead.
+
+.. warning:: C++ code that receives a raw pointer ``T *obj`` *must not*
+   assume that it has exclusive ownership of ``obj``, or even that
+   ``obj`` is allocated on the C++ heap (via ``operator new``);
+   ``obj`` might instead be a subobject of a nanobind instance
+   allocated from Python. This applies even if ``T`` supports
+   ``shared_from_this()`` and there is no associated
+   ``std::shared_ptr``. Lack of a ``shared_ptr`` does *not* imply
+   exclusive ownership; it just means there's no way to share ownership
+   with whoever the current owner is.
 
 .. _unique_ptr_adv:
 

docs/porting.rst

@@ -132,16 +132,22 @@ both of the following include directives to your code:
 .. code-block:: cpp
 
    #include <nanobind/stl/unique_ptr.h>
-   #include <nanobind/stl/unique_shared_ptr.h>
+   #include <nanobind/stl/shared_ptr.h>
 
 Binding functions that take ``std::unique_ptr<T>`` arguments involves some
 limitations that can be avoided by changing their signatures to
 ``std::unique_ptr<T, nb::deleter<T>>`` (:ref:`details <unique_ptr>`).
 
-Usage of ``std::enable_shared_from_this<T>`` is **prohibited** and will raise a
-compile-time assertion (:ref:`details <enable_shared_from_this>`) . This is
-consistent with the philosophy of this library: *the codebase has to adapt to
-the binding tool and not the other way around*.
+Use of ``std::enable_shared_from_this<T>`` is permitted, but since
+nanobind does not use holder types, an object
+constructed in Python will typically not have any associated
+``std::shared_ptr<T>`` until it is passed to a C++ function that
+accepts ``std::shared_ptr<T>``. That means a C++ function that accepts
+a raw ``T*`` and calls ``shared_from_this()`` on it might stop working
+when ported from pybind11 to nanobind. You can solve this problem
+by always passing such objects across the Python/C++ boundary as
+``std::shared_ptr<T>`` rather than as ``T*``. See the :ref:`advanced section
+on object ownership <enable_shared_from_this>` for more details.
 
 Custom constructors
 -------------------

include/nanobind/nb_class.h

@@ -48,7 +48,11 @@ enum class type_flags : uint32_t {
     /// Is this a trampoline class meant to be overloaded in Python?
     is_trampoline            = (1 << 12),
 
-    // Six more flag bits available (13 through 18) without needing
+    /// Is this a class that inherits from enable_shared_from_this?
+    /// If so, type_data::keep_shared_from_this_alive is also set.
+    has_shared_from_this     = (1 << 13),
+
+    // Five more flag bits available (14 through 18) without needing
     // a larger reorganization
 };
 
@@ -90,6 +94,7 @@ struct type_data {
     const std::type_info **implicit;
     bool (**implicit_py)(PyTypeObject *, PyObject *, cleanup_list *) noexcept;
     void (*set_self_py)(void *, PyObject *) noexcept;
+    bool (*keep_shared_from_this_alive)(PyObject *) noexcept;
 #if defined(Py_LIMITED_API)
     size_t dictoffset;
 #endif
@@ -386,6 +391,24 @@ class class_ : public object {
             }
         }
 
+        if constexpr (detail::has_shared_from_this_v<T>) {
+            d.flags |= (uint32_t) detail::type_flags::has_shared_from_this;
+            d.keep_shared_from_this_alive = [](PyObject *self) noexcept {
+                // weak_from_this().lock() is equivalent to shared_from_this(),
+                // except that it returns an empty shared_ptr instead of
+                // throwing an exception if there is no active shared_ptr
+                // for this object. (Added in C++17.)
+                if (auto sp = inst_ptr<T>(self)->weak_from_this().lock()) {
+                    detail::keep_alive(self, new auto(std::move(sp)),
+                                       [](void *p) noexcept {
+                                           delete (decltype(sp) *) p;
+                                       });
+                    return true;
+                }
+                return false;
+            };
+        }
+
         (detail::type_extra_apply(d, extra), ...);
 
         m_ptr = detail::nb_type_new(&d);

include/nanobind/nb_traits.h

@@ -132,6 +132,16 @@ struct detector<std::void_t<Op<Arg>>, Op, Arg>
    avoid redundancy when combined with nb::arg(...).none(). */
 template <typename T> struct remove_opt_mono { using type = T; };
 
+// Detect std::enable_shared_from_this without including <memory>
+template <typename T>
+auto has_shared_from_this_impl(T *ptr) ->
+    decltype(ptr->weak_from_this().lock().get(), std::true_type{});
+std::false_type has_shared_from_this_impl(...);
+
+template <typename T>
+constexpr bool has_shared_from_this_v =
+    decltype(has_shared_from_this_impl((T *) nullptr))::value;
+
 NAMESPACE_END(detail)
 
 template <typename... Args>

include/nanobind/stl/shared_ptr.h

@@ -15,33 +15,40 @@
 NAMESPACE_BEGIN(NB_NAMESPACE)
 NAMESPACE_BEGIN(detail)
 
+// shared_ptr deleter that reduces the reference count of a Python object
+struct py_deleter {
+    void operator()(void *) noexcept {
+        // Don't run the deleter if the interpreter has been shut down
+        if (!Py_IsInitialized())
+            return;
+        gil_scoped_acquire guard;
+        Py_DECREF(o);
+    }
+
+    PyObject *o;
+};
+
 /**
- * Create a generic std::shared_ptr to evade population of a potential
- * std::enable_shared_from_this weak pointer. The specified deleter reduces the
- * reference count of the Python object.
+ * Create a std::shared_ptr for `ptr` that owns a reference to the Python
+ * object `h`; if `ptr` is non-null, then the refcount of `h` is incremented
+ * before creating the shared_ptr and decremented by its deleter.
+ *
+ * Usually this is instantiated with T = void, to reduce template bloat.
+ * But if the pointee type uses enable_shared_from_this, we instantiate
+ * with T = that type, in order to allow its internal weak_ptr to share
+ * ownership with the shared_ptr we're creating.
  *
  * The next two functions are simultaneously marked as 'inline' (to avoid
  * linker errors) and 'NB_NOINLINE' (to avoid them being inlined into every
  * single shared_ptr type_caster, which would enlarge the binding size)
  */
-inline NB_NOINLINE std::shared_ptr<void>
-shared_from_python(void *ptr, handle h) noexcept {
-    struct py_deleter {
-        void operator()(void *) noexcept {
-            // Don't run the deleter if the interpreter has been shut down
-            if (!Py_IsInitialized())
-                return;
-            gil_scoped_acquire guard;
-            Py_DECREF(o);
-        }
-
-        PyObject *o;
-    };
-
+template <typename T>
+inline NB_NOINLINE std::shared_ptr<T>
+shared_from_python(T *ptr, handle h) noexcept {
     if (ptr)
-        return std::shared_ptr<void>(ptr, py_deleter{ h.inc_ref().ptr() });
+        return std::shared_ptr<T>(ptr, py_deleter{ h.inc_ref().ptr() });
     else
-        return std::shared_ptr<void>((PyObject *) nullptr);
+        return std::shared_ptr<T>(nullptr);
 }
 
 inline NB_NOINLINE void shared_from_cpp(std::shared_ptr<void> &&ptr,
@@ -50,26 +57,13 @@ inline NB_NOINLINE void shared_from_cpp(std::shared_ptr<void> &&ptr,
                [](void *p) noexcept { delete (std::shared_ptr<void> *) p; });
 }
 
-template <class T, class = void>
-struct uses_shared_from_this : std::false_type { };
-
-template <class T>
-struct uses_shared_from_this<
-    T, std::void_t<decltype(std::declval<T>().shared_from_this())>>
-    : std::true_type { };
-
 template <typename T> struct type_caster<std::shared_ptr<T>> {
     using Value = std::shared_ptr<T>;
     using Caster = make_caster<T>;
     static_assert(Caster::IsClass,
                   "Binding 'shared_ptr<T>' requires that 'T' can also be bound "
                   "by nanobind. It appears that you specified a type which "
                   "would undergo conversion/copying, which is not allowed.");
-    static_assert(!uses_shared_from_this<T>::value,
-                  "nanobind does not permit use of std::shared_from_this, "
-                  "which can cause undefined behavior. (Refer to "
-                  "https://nanobind.readthedocs.io/en/latest/ownership.html "
-                  "for details.)");
 
     static constexpr auto Name = Caster::Name;
     static constexpr bool IsClass = true;
@@ -84,9 +78,23 @@ template <typename T> struct type_caster<std::shared_ptr<T>> {
         if (!caster.from_python(src, flags, cleanup))
             return false;
 
-        value = std::static_pointer_cast<T>(
-            shared_from_python(caster.operator T *(), src));
-
+        T *ptr = caster.operator T *();
+        if constexpr (has_shared_from_this_v<T>) {
+            if (ptr) {
+                if (auto sp = ptr->weak_from_this().lock()) {
+                    // There is already a C++ shared_ptr for this object. Use it.
+                    value = std::static_pointer_cast<T>(std::move(sp));
+                    return true;
+                }
+            }
+            // Otherwise create a new one. Use shared_from_python<T>(...)
+            // so that future calls to ptr->shared_from_this() can share
+            // ownership with it.
+            value = shared_from_python(ptr, src);
+        } else {
+            value = std::static_pointer_cast<T>(
+                shared_from_python(static_cast<void *>(ptr), src));
+        }
         return true;
     }