fssystem: revise allocator logic for latest semantics

libraries/libstratosphere/include/stratosphere/fssystem/fssystem_allocator_utility.hpp

@@ -16,50 +16,100 @@
 #pragma once
 #include <vapours.hpp>
 
+/* Forward declare ams::fs allocate shared. */
+namespace ams::fs::impl {
+
+    template<typename T, template<typename, typename> class AllocatorTemplateT, typename Impl, typename... Args>
+    std::shared_ptr<T> AllocateSharedImpl(Args &&... args);
+
+}
+
 namespace ams::fssystem {
 
     using AllocateFunction   = void *(*)(size_t size);
     using DeallocateFunction = void (*)(void *ptr, size_t size);
 
     void InitializeAllocator(AllocateFunction allocate_func, DeallocateFunction deallocate_func);
+    void InitializeAllocatorForSystem(AllocateFunction allocate_func, DeallocateFunction deallocate_func);
 
     void *Allocate(size_t size);
     void Deallocate(void *ptr, size_t size);
 
-    template<typename T>
+    namespace impl {
-    class StdAllocator : public std::allocator<T> {
-        public:
-            StdAllocator() = default;
-            StdAllocator(const StdAllocator &) = default;
-            template<class U>
-            StdAllocator(const StdAllocator<U> &) : std::allocator<T>() { /* ... */ };
 
-            template<typename U>
+        template<bool ForSystem>
-            struct rebind {
+        class AllocatorFunctionSet {
-                using other = StdAllocator<U>;
+            public:
-            };
+                static void *Allocate(size_t size);
+                static void *AllocateUnsafe(size_t size);
 
-            T *Allocate(size_t size, const T *hint = nullptr) {
+                static void Deallocate(void *ptr, size_t size);
-                AMS_UNUSED(hint);
+                static void DeallocateUnsafe(void *ptr, size_t size);
-                return static_cast<T *>(::ams::fssystem::Allocate(sizeof(T) * size));
-            }
 
-            void Deallocate(T *ptr, size_t size) {
+                static void LockAllocatorMutex();
-                return ::ams::fssystem::Deallocate(ptr, sizeof(T) * size);
+                static void UnlockAllocatorMutex();
-            }
+        };
 
-            ALWAYS_INLINE T *allocate(size_t size, const T *hint = nullptr) { return this->Allocate(size, hint); }
+        using AllocatorFunctionSetForNormal = AllocatorFunctionSet<false>;
-            ALWAYS_INLINE void deallocate(T *ptr, size_t size) { return this->Deallocate(ptr, size); }
+        using AllocatorFunctionSetForSystem = AllocatorFunctionSet<true>;
-    };
+
+        template<typename T, typename Impl, bool RequireNonNull, bool AllocateWhileLocked>
+        class AllocatorTemplate : public std::allocator<T> {
+            public:
+                template<typename U>
+                struct rebind {
+                    using other = AllocatorTemplate<U, Impl, RequireNonNull, AllocateWhileLocked>;
+                };
+            private:
+                static ALWAYS_INLINE T *AllocateImpl(::std::size_t n) {
+                    if constexpr (AllocateWhileLocked) {
+                        auto * const p = Impl::AllocateUnsafe(sizeof(T) * n);
+                        Impl::UnlockAllocatorMutex();
+                        return static_cast<T *>(p);
+                    } else {
+                        return static_cast<T *>(Impl::Allocate(sizeof(T) * n));
+                    }
+                }
+            public:
+                AllocatorTemplate() { /* ... */ }
+
+                template<typename U>
+                AllocatorTemplate(const AllocatorTemplate<U, Impl, RequireNonNull, AllocateWhileLocked> &) { /* ... */ }
+
+                [[nodiscard]] T *allocate(::std::size_t n) {
+                    auto * const p = AllocateImpl(n);
+                    if constexpr (RequireNonNull) {
+                        AMS_ABORT_UNLESS(p != nullptr);
+                    }
+                    return p;
+                }
+
+                void deallocate(T *p, ::std::size_t n) {
+                    Impl::Deallocate(p, sizeof(T) * n);
+                }
+        };
+
+        template<typename T, typename Impl>
+        using AllocatorTemplateForAllocateShared = AllocatorTemplate<T, Impl, true, true>;
 
-    template<typename T>
-    std::shared_ptr<T> AllocateShared() {
-        return std::allocate_shared<T>(StdAllocator<T>{});
     }
 
     template<typename T, typename... Args>
     std::shared_ptr<T> AllocateShared(Args &&... args) {
-        return std::allocate_shared<T>(StdAllocator<T>{}, std::forward<Args>(args)...);
+        return ::ams::fs::impl::AllocateSharedImpl<T, impl::AllocatorTemplateForAllocateShared, impl::AllocatorFunctionSetForNormal>(std::forward<Args>(args)...);
+    }
+
+    template<typename TImpl, typename ErrorResult, typename TIntf, typename... Args>
+    Result AllocateSharedForSystem(std::shared_ptr<TIntf> *out, Args &&... args) {
+        /* Allocate the object. */
+        auto p = ::ams::fs::impl::AllocateSharedImpl<TImpl, impl::AllocatorTemplateForAllocateShared, impl::AllocatorFunctionSetForSystem>(std::forward<Args>(args)...);
+
+        /* Check that allocation succeeded. */
+        R_UNLESS(p != nullptr, ErrorResult());
+
+        /* Return the allocated object. */
+        *out = std::move(p);
+        return ResultSuccess();
     }
 
 }

libraries/libstratosphere/source/fssystem/fssystem_allocator_utility.cpp

@@ -21,37 +21,110 @@ namespace ams::fssystem {
 
         constexpr bool UseDefaultAllocators = false;
 
-        bool g_used_default_allocator;
+        constinit bool g_used_default_allocator = false;
 
         void *DefaultAllocate(size_t size) {
             g_used_default_allocator = true;
-            return std::malloc(size);
+            return ams::Malloc(size);
         }
 
         void DefaultDeallocate(void *ptr, size_t size) {
             AMS_UNUSED(size);
-            std::free(ptr);
+            ams::Free(ptr);
         }
 
-        AllocateFunction g_allocate_func     = UseDefaultAllocators ? DefaultAllocate   : nullptr;
+        constinit os::SdkMutex g_allocate_mutex;
-        DeallocateFunction g_deallocate_func = UseDefaultAllocators ? DefaultDeallocate : nullptr;
+        constinit os::SdkMutex g_allocate_mutex_for_system;
+
+        constinit AllocateFunction g_allocate_func     = UseDefaultAllocators ? DefaultAllocate   : nullptr;
+        constinit DeallocateFunction g_deallocate_func = UseDefaultAllocators ? DefaultDeallocate : nullptr;
+
+        constinit AllocateFunction g_allocate_func_for_system     = nullptr;
+        constinit DeallocateFunction g_deallocate_func_for_system = nullptr;
+
+        void *AllocateUnsafe(size_t size) {
+            /* Check pre-conditions. */
+            AMS_ASSERT(g_allocate_mutex.IsLockedByCurrentThread());
+            AMS_ASSERT(g_allocate_func != nullptr);
+
+            /* Allocate memory. */
+            return g_allocate_func(size);
+        }
+
+        void DeallocateUnsafe(void *ptr, size_t size) {
+            /* Check pre-conditions. */
+            AMS_ASSERT(g_allocate_mutex.IsLockedByCurrentThread());
+            AMS_ASSERT(g_deallocate_func != nullptr);
+
+            /* Deallocate the pointer. */
+            g_deallocate_func(ptr, size);
+        }
+
+
+    }
+
+    namespace impl {
+
+        /* Normal allocator set. */
+        template<> void *AllocatorFunctionSet<false>::Allocate(size_t size) { return ::ams::fssystem::Allocate(size); }
+        template<> void *AllocatorFunctionSet<false>::AllocateUnsafe(size_t size)  { return ::ams::fssystem::AllocateUnsafe(size); }
+
+        template<> void AllocatorFunctionSet<false>::Deallocate(void *ptr, size_t size) { return ::ams::fssystem::Deallocate(ptr, size); }
+        template<> void AllocatorFunctionSet<false>::DeallocateUnsafe(void *ptr, size_t size) { return ::ams::fssystem::DeallocateUnsafe(ptr, size); }
+
+        template<> void AllocatorFunctionSet<false>::LockAllocatorMutex() { ::ams::fssystem::g_allocate_mutex.Lock(); }
+        template<> void AllocatorFunctionSet<false>::UnlockAllocatorMutex() { ::ams::fssystem::g_allocate_mutex.Unlock(); }
+
+        /* System allocator set. */
+        template<> void *AllocatorFunctionSet<true>::AllocateUnsafe(size_t size)  {
+            /* Check pre-conditions. */
+            AMS_ASSERT(::ams::fssystem::g_allocate_mutex_for_system.IsLockedByCurrentThread());
+            AMS_ASSERT(::ams::fssystem::g_allocate_func_for_system != nullptr);
+
+            /* Allocate memory. */
+            return g_allocate_func_for_system(size);
+        }
+
+        template<> void AllocatorFunctionSet<true>::DeallocateUnsafe(void *ptr, size_t size) {
+            /* Check pre-conditions. */
+            AMS_ASSERT(::ams::fssystem::g_allocate_mutex_for_system.IsLockedByCurrentThread());
+            AMS_ASSERT(::ams::fssystem::g_deallocate_func_for_system != nullptr);
+
+            /* Deallocate the pointer. */
+            ::ams::fssystem::g_deallocate_func_for_system(ptr, size);
+        }
+
+        template<> void *AllocatorFunctionSet<true>::Allocate(size_t size) {
+            std::scoped_lock lk(::ams::fssystem::g_allocate_mutex_for_system);
+            return ::ams::fssystem::impl::AllocatorFunctionSet<true>::AllocateUnsafe(size);
+        }
+
+        template<> void AllocatorFunctionSet<true>::Deallocate(void *ptr, size_t size) {
+            std::scoped_lock lk(::ams::fssystem::g_allocate_mutex_for_system);
+            return ::ams::fssystem::impl::AllocatorFunctionSet<true>::DeallocateUnsafe(ptr, size);
+        }
+
+        template<> void AllocatorFunctionSet<true>::LockAllocatorMutex() { ::ams::fssystem::g_allocate_mutex_for_system.Lock(); }
+        template<> void AllocatorFunctionSet<true>::UnlockAllocatorMutex() { ::ams::fssystem::g_allocate_mutex_for_system.Unlock(); }
 
     }
 
     void *Allocate(size_t size) {
-        AMS_ASSERT(g_allocate_func != nullptr);
+        std::scoped_lock lk(g_allocate_mutex);
-        return g_allocate_func(size);
+        return AllocateUnsafe(size);
     }
 
     void Deallocate(void *ptr, size_t size) {
-        AMS_ASSERT(g_deallocate_func != nullptr);
+        std::scoped_lock lk(g_allocate_mutex);
-        return g_deallocate_func(ptr, size);
+        return DeallocateUnsafe(ptr, size);
     }
 
     void InitializeAllocator(AllocateFunction allocate_func, DeallocateFunction deallocate_func) {
+        /* Check pre-conditions. */
         AMS_ASSERT(allocate_func != nullptr);
         AMS_ASSERT(deallocate_func != nullptr);
 
+        /* Check that we can initialize. */
         if constexpr (UseDefaultAllocators) {
             AMS_ASSERT(g_used_default_allocator == false);
         } else {
@@ -59,8 +132,23 @@ namespace ams::fssystem {
             AMS_ASSERT(g_deallocate_func == nullptr);
         }
 
+        /* Set the allocator functions. */
         g_allocate_func   = allocate_func;
         g_deallocate_func = deallocate_func;
     }
 
+    void InitializeAllocatorForSystem(AllocateFunction allocate_func, DeallocateFunction deallocate_func) {
+        /* Check pre-conditions. */
+        AMS_ASSERT(allocate_func != nullptr);
+        AMS_ASSERT(deallocate_func != nullptr);
+
+        /* Check that we can initialize. */
+        AMS_ASSERT(g_allocate_func_for_system == nullptr);
+        AMS_ASSERT(g_deallocate_func_for_system == nullptr);
+
+        /* Set the system allocator functions. */
+        g_allocate_func_for_system   = allocate_func;
+        g_deallocate_func_for_system = deallocate_func;
+    }
+
 }

libraries/libstratosphere/source/fssystem/fssystem_file_system_proxy_api.cpp

@@ -124,6 +124,7 @@ namespace ams::fssystem {
         /* TODO FS-REIMPL: sf::SetGlobalDefaultMemoryResource() */
         fs::SetAllocator(AllocateForFileSystemProxy, DeallocateForFileSystemProxy);
         fssystem::InitializeAllocator(AllocateForFileSystemProxy, DeallocateForFileSystemProxy);
+        fssystem::InitializeAllocatorForSystem(AllocateForFileSystemProxy, DeallocateForFileSystemProxy);
 
         /* Initialize the buffer manager. */
         /* TODO FS-REIMPL: os::AllocateMemoryBlock(...); */