Misaki.HighPerformance/Misaki.HighPerformance.LowLevel/Buffer/AllocationManager.cs

#if MHP_ENABLE_SAFETY_CHECKS
using Misaki.HighPerformance.Collections;
#endif
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace Misaki.HighPerformance.LowLevel.Buffer;

/// <summary>
/// Holds information about a memory allocation.
/// </summary>
public readonly struct AllocationInfo
{
    /// <summary>
    /// Gets the address of the allocated memory block.
    /// </summary>
    public IntPtr Address
    {
        get; init;
    }

    /// <summary>
    /// Gets the size of the allocation in bytes.
    /// </summary>
    public nuint Size
    {
        get; init;
    }

#if MHP_ENABLE_STACKTRACE
    /// <summary>
    /// Gets the stack trace at the time of allocation for debugging purposes.
    /// </summary>
    public StackTrace? StackTrace
    {
        get; init;
    }
#endif
}

public struct AllocationManagerDesc
{
    public nuint ArenaCapacity
    {
        get; set;
    }

    public nuint StackCapacity
    {
        get; set;
    }

    public nuint FreeListChunkSize
    {
        get; set;
    }

    public nuint FreeListDefaultAlignment
    {
        get; set;
    }

    public nuint TLSFAlignment
    {
        get; set;
    }

    public nuint TLSFInitialChunkSize
    {
        get; set;
    }

    public static AllocationManagerDesc Default => new AllocationManagerDesc
    {
        ArenaCapacity = 1024 * 1024 * 1024, // 1 GB
        StackCapacity = 32 * 1024 * 1024, // 32 MB per thread
        FreeListChunkSize = 64 * 1024 * 1024,
        FreeListDefaultAlignment = 8,
        TLSFAlignment = 16,
        TLSFInitialChunkSize = 64 * 1024, // 64 KB
    };
}

/// <summary>
/// Provides memory allocation management for native memory allocations, with support for tracking,
/// debugging, and custom allocation strategies.
/// </summary>
public static unsafe class AllocationManager
{
    internal struct HeapAllocator : IAllocator
    {
        private AllocationHandle _handle;

        public readonly AllocationHandle Handle => _handle;

        public void Init()
        {
            _handle = new AllocationHandle(null, &Allocate, &Reallocate, &Free);
        }

        private static void* Allocate(void* _, nuint size, nuint alignment, AllocationOption allocationOption)
        {
            var ptr = AlignedAlloc(size, alignment);
            if (ptr == null)
            {
                return null;
            }

            if (allocationOption.HasFlag(AllocationOption.Clear))
            {
                MemClear(ptr, size);
            }

            return ptr;
        }

        private static void* Reallocate(void* _, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption)
        {
            var newPtr = AlignedRealloc(ptr, newSize, alignment);
            if (newPtr == null)
            {
                return null;
            }

            if (allocationOption.HasFlag(AllocationOption.Clear) && newSize > oldSize)
            {
                var offset = (byte*)newPtr + oldSize;
                var clearSize = newSize - oldSize;
                MemClear(offset, clearSize);
            }

            return newPtr;
        }

        private static void Free(void* _, void* ptr)
        {
            AlignedFree(ptr);
        }
    }

    // TODO: Lock-free implementation
    internal struct TLSFAllocator : IAllocator, IDisposable
    {
        private TLSF _tlsf;
        private GCHandle _lock;
        private AllocationHandle _handle;

        public readonly AllocationHandle Handle => _handle;

        public void Init(nuint alignment, nuint initialChunkSize)
        {
            _tlsf = new TLSF(alignment, initialChunkSize);
#pragma warning disable CS9216 // A value of type 'System.Threading.Lock' converted to a different type will use likely unintended monitor-based locking in 'lock' statement.
            _lock = GCHandle.Alloc(new Lock(), GCHandleType.Normal);
#pragma warning restore CS9216 // A value of type 'System.Threading.Lock' converted to a different type will use likely unintended monitor-based locking in 'lock' statement.
            _handle = new AllocationHandle(Unsafe.AsPointer(in this), &Allocate, &Reallocate, &Free);
        }

        private static void* Allocate(void* state, nuint size, nuint alignment, AllocationOption allocationOption)
        {
            var allocator = (TLSFAllocator*)state;
            var locker = (Lock)allocator->_lock.Target!;

            lock (locker)
            {
                return allocator->_tlsf.Allocate(size, alignment, allocationOption);
            }
        }

        private static void* Reallocate(void* state, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption)
        {
            var allocator = (TLSFAllocator*)state;
            var locker = (Lock)allocator->_lock.Target!;

            lock (locker)
            {
                return allocator->_tlsf.Reallocate(ptr, oldSize, newSize, alignment, allocationOption);
            }
        }

        private static void Free(void* state, void* ptr)
        {
            var allocator = (TLSFAllocator*)state;
            var locker = (Lock)allocator->_lock.Target!;

            lock (locker)
            {
                allocator->_tlsf.Free(ptr);
            }
        }

        public void Dispose()
        {
            _lock.Free();
        }
    }

    private class ThreadLocalStackPool
    {
        public MemoryPool<VirtualStack, VirtualStack.CreationOptions> pool;

        public ThreadLocalStackPool(nuint stackCapacity)
        {
            pool = new MemoryPool<VirtualStack, VirtualStack.CreationOptions>(new VirtualStack.CreationOptions
            {
                reserveCapacity = stackCapacity
            });
        }

        ~ThreadLocalStackPool()
        {
            pool.Dispose();
        }
    }

    internal static MemoryPool<VirtualArena, VirtualArena.CreationOptions> s_arenaAllocator;
    internal static MemoryPool<FreeList, FreeList.CreationOptions> s_freeListAllocator;
    internal static HeapAllocator* s_pHeapAllocator;
    internal static TLSFAllocator* s_pTLSFAllocator;

    [ThreadStatic]
    private static ThreadLocalStackPool? t_stackAllocator;


#if MHP_ENABLE_SAFETY_CHECKS
    private static ConcurrentSlotMap<AllocationInfo> s_allocations = null!;
    private static long s_totalAllocatedMemory;
#endif

    /// <summary>
    /// Gets the number of live tracked heap allocations. Always returns 0 if MHP_ENABLE_SAFETY_CHECKS is disabled.
    /// </summary>
    public static int LiveAllocationCount =>
#if MHP_ENABLE_SAFETY_CHECKS
        s_allocations.Count;
#else
        0;
#endif

    public static nuint TotalAllocatedMemory =>
#if MHP_ENABLE_SAFETY_CHECKS
        (nuint)s_totalAllocatedMemory;
#else
        0;
#endif

    private static volatile bool s_initialized;

    private static nuint s_threadLocalStackSize;

    public static void Initialize(AllocationManagerDesc desc = default)
    {
        if (s_initialized)
        {
            return;
        }

#if MHP_ENABLE_SAFETY_CHECKS
        s_allocations = new ConcurrentSlotMap<AllocationInfo>(256);
#endif

        var defaultDesc = AllocationManagerDesc.Default;

        var spanDesc = MemoryMarshal.AsBytes(MemoryMarshal.CreateSpan(ref desc, 1));
        var spanDefault = MemoryMarshal.AsBytes(MemoryMarshal.CreateSpan(ref defaultDesc, 1));
        ReplaceIfZeros(spanDesc, spanDefault);

        s_arenaAllocator = new MemoryPool<VirtualArena, VirtualArena.CreationOptions>(new VirtualArena.CreationOptions
        {
            reserveCapacity = desc.ArenaCapacity
        });

        s_freeListAllocator = new MemoryPool<FreeList, FreeList.CreationOptions>(new FreeList.CreationOptions
        {
            alignment = desc.FreeListDefaultAlignment,
            chunkSize = desc.FreeListChunkSize
        });

        s_pHeapAllocator = (HeapAllocator*)Malloc((nuint)sizeof(HeapAllocator));
        s_pHeapAllocator->Init();

        s_pTLSFAllocator = (TLSFAllocator*)Malloc((nuint)sizeof(TLSFAllocator));
        s_pTLSFAllocator->Init(desc.TLSFAlignment, desc.TLSFInitialChunkSize);

        s_threadLocalStackSize = desc.StackCapacity;

        s_initialized = true;
    }

    /// <summary>
    /// Gets a reference to the allocation pHandle for the specified allocator type.
    /// </summary>
    /// <param name="allocator">The allocator type for which to retrieve the allocation pHandle.</param>
    /// <returns>A reference to the allocation pHandle associated with the specified allocator type.</returns>
    /// <exception cref="ArgumentException"></exception>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    [Obsolete("Use AllocationHandle instead.")]
    public static AllocationHandle GetAllocationHandle(Allocator allocator)
    {
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");

        return allocator switch
        {
            Allocator.Temp => s_arenaAllocator.AllocationHandle,
            Allocator.Persistent => s_pHeapAllocator->Handle,
            Allocator.FreeList => s_freeListAllocator.AllocationHandle,
            _ => throw new ArgumentException("Target allocator type does not support custom allocation.", nameof(allocator)),
        };
    }

    /// <summary>
    /// Resets the temporary memory allocator, clearing all allocated memory.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static void ResetTempAllocator()
    {
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");
        s_arenaAllocator.Allocator.Reset();
    }

    /// <summary>
    /// Creates a new thread local stack scope for managing temporary allocations within the current context.
    /// </summary>
    /// <returns>A <see cref="Stack.Scope"/> instance representing the newly created stack scope. The scope must be disposed when no longer needed to release allocated resources.</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static VirtualStack.Scope CreateStackScope()
    {
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");

        t_stackAllocator ??= new ThreadLocalStackPool(s_threadLocalStackSize);
        return t_stackAllocator.pool.Allocator.CreateScope(t_stackAllocator.pool.AllocationHandle);
    }

    /// <summary>
    /// Registers a memory allocation and returns a handle that can be used to manage or reference the allocated memory.
    /// </summary>
    /// <remarks>
    /// Always returns an invalid handle if MHP_ENABLE_SAFETY_CHECKS is disabled.
    /// </remarks>
    /// <param name="ptr">A pointer to the memory block to be registered. The pointer must reference a valid, allocated memory region.</param>
    /// <param name="size">The newSize of the memory block to be registered.</param>
    /// <returns>A MemoryHandle representing the registered allocation.</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static MemoryHandle AddAllocation(void* ptr, nuint size)
    {
#if MHP_ENABLE_SAFETY_CHECKS
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");

        var info = new AllocationInfo
        {
            Address = (IntPtr)ptr,
            Size = size,
#if MHP_ENABLE_STACKTRACE
            StackTrace = new StackTrace(1, true)
#endif
        };

        Interlocked.Add(ref s_totalAllocatedMemory, (long)size);

        var id = s_allocations.Add(info, out var generation);
        return new MemoryHandle(id, generation);
#else
        return MemoryHandle.Invalid;
#endif
    }

    public static void UpdateAllocation(MemoryHandle handle, void* newPtr, nuint newSize)
    {
#if MHP_ENABLE_SAFETY_CHECKS
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");

        if (newPtr == null)
        {
            s_allocations.Remove(handle.ID, handle.Generation, out var info);
            Interlocked.Add(ref s_totalAllocatedMemory, -(long)info.Size);
        }
        else
        {

            if (s_allocations.TryGetElement(handle.ID, handle.Generation, out var oldInfo))
            {
                var diff = newSize - oldInfo.Size;
                Interlocked.Add(ref s_totalAllocatedMemory, (long)diff);

                var newInfo = oldInfo with
                {
                    Address = (IntPtr)newPtr,
                    Size = newSize
                };

                s_allocations.UpdateElement(handle.ID, handle.Generation, newInfo);
            }
        }
#endif
    }

    /// <summary>
    /// Removes the memory allocation associated with the specified handle.
    /// </summary>
    /// <remarks>
    /// Always returns false if debug layer is disabled.
    /// </remarks>
    /// <param name="handle">The handle representing the memory allocation to remove. The handle must be valid and previously allocated.</param>
    /// <returns>true if the allocation was successfully removed; otherwise, false.</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static bool RemoveAllocation(MemoryHandle handle)
    {
#if MHP_ENABLE_SAFETY_CHECKS
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");

        if (s_allocations.Remove(handle.ID, handle.Generation, out var info))
        {
            Interlocked.Add(ref s_totalAllocatedMemory, -(long)info.Size);
            return true;
        }

        return false;
#else
        return false;
#endif
    }

    /// <summary>
    /// Attempts to retrieve the memory allocation pointer associated with the specified handle.
    /// </summary>
    /// <remarks>
    /// Always returns false if debug layer is disabled, and the output pointer will be set to <see cref="IntPtr.Zero"/>.
    /// </remarks>
    /// <param name="handle">The memory handle identifying the allocation to retrieve allocation.</param>
    /// <param name="info">When this method returns, contains the allocation information associated with the specified handle, if the allocation was found; otherwise, an uninitialized value.</param>
    /// <returns>true if the allocation was found and <paramref name="info"/> contains valid allocation information; otherwise, false.</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static bool TryGetAllocation(MemoryHandle handle, out AllocationInfo info)
    {
#if MHP_ENABLE_SAFETY_CHECKS
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");
        return s_allocations.TryGetElement(handle.ID, handle.Generation, out info);
#else
        info = default;
        return false;
#endif
    }

    /// <summary>
    /// Determines whether the specified memory handle refers to a currently tracked allocation.
    /// </summary>
    /// <remarks>
    /// This only validates the memory when you added the allocation via <see cref="AddAllocation(IntPtr)"/>.
    /// For validating memory from <see cref="AllocationHandle"/>, use <see cref="AllocationHandle.IsValid"/> instead.
    /// Always returns false if debug layer is disabled.
    /// </remarks>
    /// <param name="handle">The memory handle to check for an associated allocation.</param>
    /// <returns>true if the allocation corresponding to the handle exists; otherwise, false.</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static bool ContainsAllocation(MemoryHandle handle)
    {
#if MHP_ENABLE_SAFETY_CHECKS
        Debug.Assert(s_initialized, "AllocationManager is not initialized.");
        return s_allocations.Contains(handle.ID, handle.Generation);
#else
        return false;
#endif
    }

    /// <summary>
    /// Disposes of the AllocationManager, freeing all allocated memory and resources.
    /// </summary>
    public static void Dispose()
    {
        if (!s_initialized)
        {
            return;
        }

        s_initialized = false;

#if MHP_ENABLE_SAFETY_CHECKS
        if (s_allocations.Count > 0)
        {
            throw new MemoryLeakException(s_allocations);
        }
#endif

        s_arenaAllocator.Dispose();
        s_freeListAllocator.Dispose();

        if (s_pHeapAllocator != null)
        {
            Free(s_pHeapAllocator);
            s_pHeapAllocator = null;
        }

        if (s_pTLSFAllocator != null)
        {
            s_pTLSFAllocator->Dispose();
            Free(s_pTLSFAllocator);
            s_pTLSFAllocator = null;
        }
    }
}