From ef62730ae457bb83b0424b11f764ee07eb0d8432 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Steven=20Dee=20=28J=C5=8Dshin=29?= <git@wholezero.org>
Date: Sun, 15 Sep 2024 19:32:13 -0400
Subject: [PATCH] Enable STL-style enable_shared_from_this (#1295)

---
 ctl/shared_ptr.h            | 65 ++++++++++++++++++++++++++++++++++---
 test/ctl/shared_ptr_test.cc | 35 ++++++++++++++++++++
 2 files changed, 96 insertions(+), 4 deletions(-)

diff --git a/ctl/shared_ptr.h b/ctl/shared_ptr.h
index 317374bfd..8bb892943 100644
--- a/ctl/shared_ptr.h
+++ b/ctl/shared_ptr.h
@@ -4,6 +4,7 @@
 #define CTL_SHARED_PTR_H_
 
 #include "exception.h"
+#include "is_base_of.h"
 #include "is_convertible.h"
 #include "remove_extent.h"
 #include "unique_ptr.h"
@@ -201,10 +202,7 @@ class shared_ptr
 
     template<typename U, typename D>
         requires __::shared_ptr_compatible<T, U>
-    shared_ptr(U* const p, D d)
-      : p(p), rc(__::shared_pointer<U, D>::make(p, move(d)))
-    {
-    }
+    shared_ptr(U*, D);
 
     template<typename U>
     shared_ptr(const shared_ptr<U>& r, element_type* p) noexcept
@@ -443,6 +441,62 @@ class weak_ptr
     __::shared_ref* rc = nullptr;
 };
 
+template<typename T>
+class enable_shared_from_this
+{
+  public:
+    shared_ptr<T> shared_from_this()
+    {
+        return shared_ptr<T>(weak_this);
+    }
+    shared_ptr<T const> shared_from_this() const
+    {
+        return shared_ptr<T>(weak_this);
+    }
+
+    weak_ptr<T> weak_from_this()
+    {
+        return weak_this;
+    }
+    weak_ptr<T const> weak_from_this() const
+    {
+        return weak_this;
+    }
+
+  protected:
+    constexpr enable_shared_from_this() noexcept = default;
+    enable_shared_from_this(const enable_shared_from_this& r) noexcept
+    {
+    }
+    ~enable_shared_from_this() = default;
+
+    enable_shared_from_this& operator=(
+      const enable_shared_from_this& r) noexcept
+    {
+        return *this;
+    }
+
+  private:
+    template<typename U, typename... Args>
+    friend shared_ptr<U> make_shared(Args&&...);
+
+    template<typename U>
+    friend class shared_ptr;
+
+    weak_ptr<T> weak_this;
+};
+
+template<typename T>
+template<typename U, typename D>
+    requires __::shared_ptr_compatible<T, U>
+shared_ptr<T>::shared_ptr(U* const p, D d)
+  : p(p), rc(__::shared_pointer<U, D>::make(p, move(d)))
+{
+    if constexpr (is_base_of_v<enable_shared_from_this<U>, U>) {
+        p->weak_this = *this;
+    }
+}
+
 template<typename T, typename... Args>
 shared_ptr<T>
 make_shared(Args&&... args)
@@ -452,6 +506,9 @@ make_shared(Args&&... args)
     shared_ptr<T> r;
     r.p = &rc->t;
     r.rc = rc.release();
+    if constexpr (is_base_of_v<enable_shared_from_this<T>, T>) {
+        r->weak_this = r;
+    }
     return r;
 }
 
diff --git a/test/ctl/shared_ptr_test.cc b/test/ctl/shared_ptr_test.cc
index 960d6b5fc..c910ddf70 100644
--- a/test/ctl/shared_ptr_test.cc
+++ b/test/ctl/shared_ptr_test.cc
@@ -25,6 +25,7 @@
 // #define ctl std
 
 using ctl::bad_weak_ptr;
+using ctl::enable_shared_from_this;
 using ctl::make_shared;
 using ctl::move;
 using ctl::shared_ptr;
@@ -66,6 +67,27 @@ struct Base
 struct Derived : Base
 {};
 
+class SharedThis : public enable_shared_from_this<SharedThis>
+{
+    struct Private
+    {
+        explicit Private() = default;
+    };
+
+  public:
+    SharedThis(Private)
+    {
+    }
+
+    static shared_ptr<SharedThis> create()
+    {
+        return make_shared<SharedThis>(Private());
+    }
+};
+
+class CanShareThis : public enable_shared_from_this<CanShareThis>
+{};
+
 int
 main()
 {
@@ -241,6 +263,19 @@ main()
             return 23;
     }
 
+    {
+        // enable_shared_from_this allows shared pointers to self.
+        auto x = SharedThis::create();
+        auto y = x->shared_from_this();
+        if (x.use_count() != 2 || x.get() != y.get())
+            return 24;
+        auto z = new CanShareThis();
+        auto w = shared_ptr<CanShareThis>(z);
+        auto v = w->shared_from_this();
+        if (w.use_count() != 2 || w.get() != v.get())
+            return 25;
+    }
+
     // TODO(mrdomino): exercise threads / races. The reference count should be
     // atomically maintained.