libluna: Store SharedPtr's ref count in the object itself

kernel/src/arch/x86_64/disk/ATA.h

@@ -131,7 +131,7 @@ namespace ATA
 
     static constexpr u16 END_OF_PRDT = (1 << 15);
 
-    class Drive
+    class Drive : public Shareable
     {
       public:
         Drive(Channel* channel, u8 drive_index, Badge<Channel>);
@@ -274,7 +274,7 @@ namespace ATA
         SharedPtr<Drive> m_drives[2];
     };
 
-    class Controller
+    class Controller : public Shareable
     {
       public:
         static Result<void> scan();

kernel/src/fs/Pipe.h

@@ -6,7 +6,7 @@ class PipeInodeBase;
 class PipeReader;
 class PipeWriter;
 
-class Pipe
+class Pipe : public Shareable
 {
   public:
     static Result<void> create(SharedPtr<VFS::Inode>& rpipe, SharedPtr<VFS::Inode>& wpipe);

kernel/src/fs/VFS.h

@@ -21,7 +21,7 @@ namespace VFS
 
     class Inode;
 
-    class FileSystem
+    class FileSystem : public Shareable
     {
       public:
         virtual SharedPtr<Inode> root_inode() const = 0;
@@ -73,7 +73,7 @@ namespace VFS
         StaticString<128> name;
     };
 
-    class Inode
+    class Inode : public Shareable
     {
       public:
         virtual Result<u64> ioctl(int, void*)

kernel/src/fs/devices/Device.h

@@ -1,8 +1,9 @@
 #pragma once
 #include "Log.h"
 #include <luna/Result.h>
+#include <luna/SharedPtr.h>
 
-class Device
+class Device : public Shareable
 {
   public:
     virtual Result<usize> read(u8* buf, usize offset, usize length) const = 0;

libluna/include/luna/SharedPtr.h

@@ -5,89 +5,76 @@
 #include <luna/OwnedPtr.h>
 #include <luna/Result.h>
 #include <luna/ScopeGuard.h>
+#include <luna/TypeTraits.h>
 
-namespace __detail
+template <typename T> class SharedPtr;
+template <typename T> SharedPtr<T> adopt_shared(T*);
+
+struct Shareable
 {
-    struct RefCount
+    void ref()
     {
-        void ref()
-        {
-            m_ref_count++;
-        }
+        m_ref_count++;
+    }
 
-        bool unref()
-        {
-            m_ref_count--;
-            return m_ref_count == 0;
-        }
+    bool unref()
+    {
+        m_ref_count--;
+        return m_ref_count == 0;
+    }
 
-      private:
-        Atomic<int> m_ref_count { 1 };
-    };
-}
+    Atomic<int> m_ref_count { 1 };
+};
 
 template <typename T> class SharedPtr
 {
-    using RefCount = __detail::RefCount;
-
   public:
     SharedPtr()
     {
         m_ptr = nullptr;
-        m_ref_count = nullptr;
     }
 
-    SharedPtr(T* ptr, RefCount* ref_count) : m_ptr(ptr), m_ref_count(ref_count)
+    SharedPtr(T* ptr) : m_ptr(ptr)
     {
     }
 
-    SharedPtr(const SharedPtr<T>& other) : m_ptr(other.m_ptr), m_ref_count(other.m_ref_count)
+    SharedPtr(const SharedPtr<T>& other) : m_ptr(other.m_ptr)
     {
-        if (m_ref_count) m_ref_count->ref();
+        if (m_ptr) shareable()->ref();
     }
 
-    SharedPtr(SharedPtr<T>&& other) : m_ptr(other.m_ptr), m_ref_count(other.m_ref_count)
+    SharedPtr(SharedPtr<T>&& other) : m_ptr(other.m_ptr)
     {
         other.m_ptr = nullptr;
-        other.m_ref_count = nullptr;
     }
 
     template <typename Tp> operator SharedPtr<Tp>()
     {
-        if (m_ref_count) m_ref_count->ref();
-        return { (Tp*)m_ptr, m_ref_count };
+        if (m_ptr) shareable()->ref();
+        return { (Tp*)m_ptr };
     }
 
     ~SharedPtr()
     {
-        if (m_ref_count && m_ref_count->unref())
-        {
-            delete m_ref_count;
-            delete m_ptr;
-        }
+        if (m_ptr && shareable()->unref()) { delete m_ptr; }
     }
 
     SharedPtr<T>& operator=(const SharedPtr<T>& other)
     {
         if (&other == this) return *this;
 
-        if (m_ref_count && m_ref_count->unref())
-        {
-            delete m_ref_count;
-            delete m_ptr;
-        }
+        if (m_ptr && shareable()->unref()) { delete m_ptr; }
 
         m_ptr = other.m_ptr;
-        m_ref_count = other.m_ref_count;
 
-        if (m_ref_count) m_ref_count->ref();
+        if (m_ptr) shareable()->ref();
 
         return *this;
     }
 
     bool operator==(const SharedPtr<T>& other)
     {
-        return m_ptr == other.m_ptr && m_ref_count == other.m_ref_count;
+        return m_ptr == other.m_ptr;
     }
 
     T* ptr() const
@@ -112,21 +99,17 @@ template <typename T> class SharedPtr
 
   private:
     T* m_ptr;
-    RefCount* m_ref_count;
+
+    Shareable* shareable()
+    {
+        static_assert(IsBaseOf<Shareable, T>);
+        return (Shareable*)m_ptr;
+    }
 };
 
-// NOTE: ptr is deleted if any of the adopt_shared* functions fail to construct a SharedPtr.
-template <typename T> Result<SharedPtr<T>> adopt_shared(T* ptr)
+template <typename T> SharedPtr<T> adopt_shared(T* ptr)
 {
-    using RefCount = __detail::RefCount;
-
-    auto guard = make_scope_guard([ptr] { delete ptr; });
-
-    RefCount* const ref_count = TRY(make<RefCount>());
-
-    guard.deactivate();
-
-    return SharedPtr<T> { ptr, ref_count };
+    return SharedPtr<T> { ptr };
 }
 
 template <typename T, class... Args> Result<SharedPtr<T>> make_shared(Args... args)
@@ -142,12 +125,10 @@ template <typename T> Result<SharedPtr<T>> adopt_shared_if_nonnull(T* ptr)
         return err(ENOMEM);
 }
 
-template <typename T> Result<SharedPtr<T>> adopt_shared_from_owned(OwnedPtr<T>&& other)
+template <typename T> SharedPtr<T> adopt_shared_from_owned(OwnedPtr<T>&& other)
 {
     T* ptr = other.m_ptr;
     other.m_ptr = nullptr;
 
-    const SharedPtr<T> shared_ptr = TRY(adopt_shared(ptr));
-
-    return shared_ptr;
+    return SharedPtr<T> { ptr };
 }

libos/include/os/Directory.h

@@ -19,7 +19,7 @@ namespace os
     /**
      * @brief An object-oriented directory handle, which is closed when all references to it go out of scope.
      */
-    class Directory
+    class Directory : public Shareable
     {
       public:
         /**

libos/include/os/File.h

@@ -22,7 +22,7 @@ namespace os
     /**
      * @brief An object-oriented file handle, which is closed when all references to it go out of scope.
      */
-    class File
+    class File : public Shareable
     {
       public:
         /**

tests/CMakeLists.txt

@@ -5,7 +5,7 @@ target_include_directories(test PUBLIC ${LUNA_BASE}/usr/include)
 
 function(luna_test SOURCE_FILE APP_NAME)
     add_executable(${APP_NAME} ${SOURCE_FILE})
-    target_compile_options(${APP_NAME} PRIVATE -Os ${COMMON_FLAGS} -Wno-write-strings)
+    target_compile_options(${APP_NAME} PRIVATE -O0 ${COMMON_FLAGS} -fno-threadsafe-statics -Wno-write-strings)
     add_dependencies(${APP_NAME} libc)
     target_include_directories(${APP_NAME} PRIVATE ${LUNA_BASE}/usr/include)
     target_link_libraries(${APP_NAME} PRIVATE test os)
@@ -19,6 +19,7 @@ luna_test(libluna/TestUtf8.cpp TestUtf8)
 luna_test(libluna/TestFormat.cpp TestFormat)
 luna_test(libluna/TestHashTable.cpp TestHashTable)
 luna_test(libluna/TestCPath.cpp TestCPath)
+luna_test(libluna/TestSharedPtr.cpp TestSharedPtr)
 luna_test(libc/TestScanf.cpp TestScanf)
 luna_test(libc/TestString.cpp TestString)
 luna_test(libc/TestEnv.cpp TestEnv)

tests/libluna/TestSharedPtr.cpp

@@ -0,0 +1,61 @@
+#include <luna/SharedPtr.h>
+#include <test.h>
+
+class SampleClass : public Shareable
+{
+  public:
+    SampleClass(int* ptr, int value) : m_ptr(ptr), m_value(value)
+    {
+    }
+
+    ~SampleClass()
+    {
+        if (m_ptr) *m_ptr = m_value;
+    }
+
+  private:
+    int* m_ptr;
+    int m_value;
+};
+
+TestResult test_shared_ptr_destroyed_on_scope_exit()
+{
+    static int value = 0;
+
+    {
+        SharedPtr<SampleClass> sample = TRY(make_shared<SampleClass>(&value, 1));
+    }
+
+    validate(value == 1);
+
+    test_success;
+}
+
+TestResult test_shared_ptr_preserved_if_at_least_one_copy_exists()
+{
+    static int value = 0;
+    static SharedPtr<SampleClass> ptr = {};
+
+    {
+        SharedPtr<SampleClass> sample = TRY(make_shared<SampleClass>(&value, 1));
+        ptr = sample;
+    }
+
+    validate(value == 0);
+
+    ptr = {};
+
+    validate(value == 1);
+
+    test_success;
+}
+
+Result<void> test_main()
+{
+    test_prelude;
+
+    run_test(test_shared_ptr_destroyed_on_scope_exit);
+    run_test(test_shared_ptr_preserved_if_at_least_one_copy_exists);
+
+    return {};
+}