From 2a94c687ea71947de7dddbeae1036a4ddd390123 Mon Sep 17 00:00:00 2001
From: "Ralf W. Grosse-Kunstleve" <rwgk@google.com>
Date: Tue, 5 Jan 2021 16:25:25 -0800
Subject: [PATCH] Adding from/as unique_ptr<T> and unique_ptr<T, D>.

---
 include/pybind11/smart_holder_poc.h | 139 +++++++++++++++++++++-------
 tests/test_smart_holder_poc.cpp     |  49 ++++++++--
 2 files changed, 149 insertions(+), 39 deletions(-)

diff --git a/include/pybind11/smart_holder_poc.h b/include/pybind11/smart_holder_poc.h
index 43759ed2d..6ea0aa24d 100644
--- a/include/pybind11/smart_holder_poc.h
+++ b/include/pybind11/smart_holder_poc.h
@@ -5,26 +5,42 @@
 
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 
+template <typename T>
+struct guarded_builtin_delete {
+  bool* flag_ptr;
+  explicit guarded_builtin_delete(bool* flag_ptr_) : flag_ptr{flag_ptr_} {}
+  void operator()(T* raw_ptr) {
+    if (*flag_ptr) delete raw_ptr;
+  }
+};
+
+template <typename T, typename D>
+struct guarded_custom_deleter {
+  bool* flag_ptr;
+  explicit guarded_custom_deleter(bool* flag_ptr_) : flag_ptr{flag_ptr_} {}
+  void operator()(T* raw_ptr) {
+    if (*flag_ptr) D()(raw_ptr);
+  }
+};
+
 struct smart_holder {
   std::shared_ptr<void> vptr;
   const std::type_info* rtti_held;
   const std::type_info* rtti_uqp_del;
   bool vptr_deleter_flag;
 
-  template <typename T>
-  struct vptr_deleter {
-    bool* flag_ptr;
-    explicit vptr_deleter(bool* flag_ptr_) : flag_ptr{flag_ptr_} {}
-    void operator()(T* raw_ptr) {
-      if (*flag_ptr) delete raw_ptr;
-    }
-  };
+  void clear() {
+    vptr.reset();
+    vptr_deleter_flag = false;
+    rtti_held = nullptr;
+    rtti_uqp_del = nullptr;
+  }
 
   smart_holder()
       : rtti_held{nullptr}, rtti_uqp_del{nullptr}, vptr_deleter_flag{false} {}
 
   template <typename T>
-  void ensure_compatible_rtti(const char* context) {
+  void ensure_compatible_rtti_held(const char* context) {
     const std::type_info* rtti_requested = &typeid(T);
     if (!(*rtti_requested == *rtti_held)) {
       throw std::runtime_error(std::string("Incompatible RTTI (") + context +
@@ -32,10 +48,12 @@ struct smart_holder {
     }
   }
 
-  void ensure_use_count_1(const char* context) {
-    if (vptr.use_count() != 1) {
-      throw std::runtime_error(std::string("Cannot disown use_count != 1 (") +
-                               context + ").");
+  template <typename D>
+  void ensure_compatible_rtti_uqp_del(const char* context) {
+    const std::type_info* rtti_requested = &typeid(D);
+    if (!(*rtti_requested == *rtti_uqp_del)) {
+      throw std::runtime_error(
+          std::string("Incompatible unique_ptr deleter (") + context + ").");
     }
   }
 
@@ -46,32 +64,89 @@ struct smart_holder {
     }
   }
 
-  template <typename T>
-  void from_raw_ptr_owned(T* raw_ptr) {
-    vptr_deleter_flag = true;
-    vptr.reset(raw_ptr, vptr_deleter<T>(&vptr_deleter_flag));
-    rtti_held = &typeid(T);
-  }
-
-  template <typename T>
-  T* as_raw_ptr_owned() {
-    static const char* context = "as_raw_ptr_owned";
-    ensure_compatible_rtti<T>(context);
-    ensure_use_count_1(context);
-    ensure_vptr_deleter_flag_true(context);
-    std::shared_ptr<T> tptr = std::static_pointer_cast<T>(vptr);
-    vptr.reset();
-    T* result = tptr.get();
-    vptr_deleter_flag = false;
-    return result;
+  void ensure_use_count_1(const char* context) {
+    if (vptr.use_count() != 1) {
+      throw std::runtime_error(std::string("Cannot disown use_count != 1 (") +
+                               context + ").");
+    }
   }
 
   template <typename T>
   std::shared_ptr<T> as_shared_ptr() {
     static const char* context = "as_shared_ptr";
-    ensure_compatible_rtti<T>(context);
+    ensure_compatible_rtti_held<T>(context);
     return std::static_pointer_cast<T>(vptr);
   }
+
+  template <typename T>
+  void from_raw_ptr_owned(T* raw_ptr) {
+    clear();
+    rtti_held = &typeid(T);
+    vptr_deleter_flag = true;
+    vptr.reset(raw_ptr, guarded_builtin_delete<T>(&vptr_deleter_flag));
+  }
+
+  template <typename T>
+  void from_raw_ptr_unowned(T* raw_ptr) {
+    clear();
+    rtti_held = &typeid(T);
+    vptr_deleter_flag = false;
+    vptr.reset(raw_ptr, guarded_builtin_delete<T>(&vptr_deleter_flag));
+  }
+
+  template <typename T>
+  T* as_raw_ptr_owned(const char* context = "as_raw_ptr_owned") {
+    ensure_compatible_rtti_held<T>(context);
+    ensure_vptr_deleter_flag_true(context);
+    ensure_use_count_1(context);
+    T* raw_ptr = static_cast<T*>(vptr.get());
+    vptr_deleter_flag = false;
+    vptr.reset();
+    return raw_ptr;
+  }
+
+  template <typename T>
+  T* as_raw_ptr_unowned() {
+    static const char* context = "as_raw_ptr_unowned";
+    ensure_compatible_rtti_held<T>(context);
+    return static_cast<T*>(vptr.get());
+  }
+
+  template <typename T>
+  void from_unique_ptr(std::unique_ptr<T>&& unq_ptr) {
+    clear();
+    rtti_held = &typeid(T);
+    vptr_deleter_flag = true;
+    vptr.reset(unq_ptr.get(), guarded_builtin_delete<T>(&vptr_deleter_flag));
+    unq_ptr.release();
+  }
+
+  template <typename T>
+  std::unique_ptr<T> as_unique_ptr() {
+    return std::unique_ptr<T>(as_raw_ptr_owned<T>("as_unique_ptr"));
+  }
+
+  template <typename T, typename D>
+  void from_unique_ptr_with_deleter(std::unique_ptr<T, D>&& unq_ptr) {
+    clear();
+    rtti_held = &typeid(T);
+    rtti_uqp_del = &typeid(D);
+    vptr_deleter_flag = true;
+    vptr.reset(unq_ptr.get(), guarded_custom_deleter<T, D>(&vptr_deleter_flag));
+    unq_ptr.release();
+  }
+
+  template <typename T, typename D>
+  std::unique_ptr<T, D> as_unique_ptr_with_deleter() {
+    static const char* context = "as_unique_ptr_with_deleter";
+    ensure_compatible_rtti_held<T>(context);
+    ensure_compatible_rtti_uqp_del<D>(context);
+    ensure_use_count_1(context);
+    T* raw_ptr = static_cast<T*>(vptr.get());
+    vptr_deleter_flag = false;
+    vptr.reset();
+    return std::unique_ptr<T, D>(raw_ptr);
+  }
 };
 
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
diff --git a/tests/test_smart_holder_poc.cpp b/tests/test_smart_holder_poc.cpp
index 09db276bb..b26363b37 100644
--- a/tests/test_smart_holder_poc.cpp
+++ b/tests/test_smart_holder_poc.cpp
@@ -10,19 +10,54 @@ namespace smart_holder_poc {
 
 inline void to_cout(std::string msg) { std::cout << msg << std::endl; }
 
+template <typename T>
+struct functor_builtin_delete {
+  void operator()(T* ptr) { delete ptr; }
+};
+
 inline void exercise() {
   to_cout("");
   namespace py = pybind11;
-  py::smart_holder hld;
-  hld.from_raw_ptr_owned(new int(13));
-  to_cout(hld.rtti_held->name());
   {
-    std::shared_ptr<int> val = hld.as_shared_ptr<int>();
-    to_cout(std::to_string(*val));
+    py::smart_holder hld;
+    hld.from_raw_ptr_owned(new int(13));
+    to_cout(hld.rtti_held->name());
+    {
+      std::shared_ptr<int> val = hld.as_shared_ptr<int>();
+      to_cout(std::to_string(*val));
+    }
+    {
+      std::unique_ptr<int> val(hld.as_raw_ptr_owned<int>());
+      to_cout(std::to_string(*val));
+    }
+  }  // namespace ;
+  {
+    std::unique_ptr<int> val(new int(13));
+    py::smart_holder hld;
+    hld.from_raw_ptr_unowned(val.get());
+    to_cout(std::to_string(*hld.as_raw_ptr_unowned<int>()));
   }
   {
-    std::unique_ptr<int> val(hld.as_raw_ptr_owned<int>());
-    to_cout(std::to_string(*val));
+    std::unique_ptr<int> val(new int(13));
+    py::smart_holder hld;
+    hld.from_unique_ptr(std::move(val));
+    to_cout(std::to_string(*hld.as_raw_ptr_unowned<int>()));
+  }
+  {
+    py::smart_holder hld;
+    hld.from_raw_ptr_owned(new int(13));
+    to_cout(std::to_string(*hld.as_unique_ptr<int>()));
+  }
+  {
+    std::unique_ptr<int, functor_builtin_delete<int>> unq_ptr(new int(13));
+    py::smart_holder hld;
+    hld.from_unique_ptr_with_deleter(std::move(unq_ptr));
+    to_cout(std::to_string(unq_ptr.get() == nullptr));
+    to_cout(std::to_string(*hld.as_raw_ptr_unowned<int>()));
+    auto unq_ptr_retrieved =
+        hld.as_unique_ptr_with_deleter<int, functor_builtin_delete<int>>();
+    to_cout(std::to_string(hld.vptr.get() == nullptr));
+    to_cout(std::to_string(*unq_ptr_retrieved));
   }
 }