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>
-    class StdAllocator : public std::allocator<T> {
-        public:
-            StdAllocator() = default;
-            StdAllocator(const StdAllocator &) = default;
-            template<class U>
-            StdAllocator(const StdAllocator<U> &) : std::allocator<T>() { /* ... */ };
+    namespace impl {
 
-            template<typename U>
-            struct rebind {
-                using other = StdAllocator<U>;
-            };
+        template<bool ForSystem>
+        class AllocatorFunctionSet {
+            public:
+                static void *Allocate(size_t size);
+                static void *AllocateUnsafe(size_t size);
 
-            T *Allocate(size_t size, const T *hint = nullptr) {
-                AMS_UNUSED(hint);
-                return static_cast<T *>(::ams::fssystem::Allocate(sizeof(T) * size));
-            }
+                static void Deallocate(void *ptr, size_t size);
+                static void DeallocateUnsafe(void *ptr, size_t size);
 
-            void Deallocate(T *ptr, size_t size) {
-                return ::ams::fssystem::Deallocate(ptr, sizeof(T) * size);
-            }
+                static void LockAllocatorMutex();
+                static void UnlockAllocatorMutex();
+        };
 
-            ALWAYS_INLINE T *allocate(size_t size, const T *hint = nullptr) { return this->Allocate(size, hint); }
-            ALWAYS_INLINE void deallocate(T *ptr, size_t size) { return this->Deallocate(ptr, size); }
-    };
+        using AllocatorFunctionSetForNormal = AllocatorFunctionSet<false>;
+        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();
     }
 
 }