fs: implement newer allocator templates

libraries/libstratosphere/include/stratosphere/fs/fs_memory_management.hpp

@@ -28,6 +28,94 @@ namespace ams::fs {
         void *Allocate(size_t size);
         void Deallocate(void *ptr, size_t size);
 
+        void LockAllocatorMutex();
+        void UnlockAllocatorMutex();
+
+        void *AllocateUnsafe(size_t size);
+        void DeallocateUnsafe(void *ptr, size_t size);
+
+        class AllocatorImpl {
+            public:
+                static ALWAYS_INLINE void *Allocate(size_t size) { return ::ams::fs::impl::Allocate(size); }
+                static ALWAYS_INLINE void *AllocateUnsafe(size_t size) { return ::ams::fs::impl::AllocateUnsafe(size); }
+
+                static ALWAYS_INLINE void Deallocate(void *ptr, size_t size) { return ::ams::fs::impl::Deallocate(ptr, size); }
+                static ALWAYS_INLINE void DeallocateUnsafe(void *ptr, size_t size) { return ::ams::fs::impl::DeallocateUnsafe(ptr, size); }
+
+                static ALWAYS_INLINE void LockAllocatorMutex() { return ::ams::fs::impl::LockAllocatorMutex(); }
+                static ALWAYS_INLINE void UnlockAllocatorMutex() { return ::ams::fs::impl::UnlockAllocatorMutex(); }
+        };
+
+        template<typename T, typename Impl, bool AllocateWhileLocked>
+        class AllocatorTemplate : public std::allocator<T> {
+            public:
+                template<typename U>
+                struct rebind {
+                    using other = AllocatorTemplate<U, Impl, AllocateWhileLocked>;
+                };
+            private:
+                bool m_allocation_failed;
+            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() : m_allocation_failed(false) { /* ... */ }
+
+                template<typename U>
+                AllocatorTemplate(const AllocatorTemplate<U, Impl, AllocateWhileLocked> &rhs) : m_allocation_failed(rhs.IsAllocationFailed()) { /* ... */ }
+
+                bool IsAllocationFailed() const { return m_allocation_failed; }
+
+                [[nodiscard]] T *allocate(::std::size_t n) {
+                    auto * const p = AllocateImpl(n);
+                    if (AMS_UNLIKELY(p == nullptr) && n) {
+                        m_allocation_failed = true;
+                    }
+                    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>;
+
+        template<typename T, template<typename, typename> class AllocatorTemplateT, typename Impl, typename... Args>
+        std::shared_ptr<T> AllocateSharedImpl(Args &&... args) {
+            /* Try to allocate. */
+            {
+                /* Acquire exclusive access to the allocator. */
+                Impl::LockAllocatorMutex();
+
+                /* Check that we can allocate memory (using overestimate of 0x80 + sizeof(T)). */
+                if (auto * const p = Impl::Allocate(0x80 + sizeof(T)); AMS_LIKELY(p != nullptr)) {
+                    /* Free the memory we allocated. */
+                    Impl::Deallocate(p, 0x80 + sizeof(T));
+
+                    /* Get allocator type. */
+                    using AllocatorType = AllocatorTemplateT<T, Impl>;
+
+                    /* Allocate the shared pointer. */
+                    return std::allocate_shared<T>(AllocatorType{}, std::forward<Args>(args)...);
+                } else {
+                    /* We can't allocate. */
+                    Impl::UnlockAllocatorMutex();
+                }
+            }
+
+            /* We failed. */
+            return nullptr;
+        }
+
         class Deleter {
             private:
                 size_t m_size;
@@ -59,4 +147,9 @@ namespace ams::fs {
 
     }
 
+    template<typename T, typename... Args>
+    std::shared_ptr<T> AllocateShared(Args &&... args) {
+        return ::ams::fs::impl::AllocateSharedImpl<T, ::ams::fs::impl::AllocatorTemplateForAllocateShared, ::ams::fs::impl::AllocatorImpl>(std::forward<Args>(args)...);
+    }
+
 }