#if MHP_ENABLE_SAFETY_CHECKS
using Misaki.HighPerformance.Collections;
using System.Collections.Concurrent;
#endif
using System.Diagnostics;
using System.Runtime.CompilerServices;
namespace Misaki.HighPerformance.LowLevel.Buffer;
///
/// Holds information about a memory allocation.
///
public readonly struct AllocationInfo
{
///
/// Gets the address of the allocated memory block.
///
public IntPtr Address
{
get; init;
}
///
/// Gets the newSize of the allocation in bytes.
///
public nuint Size
{
get; init;
}
#if MHP_ENABLE_STACKTRACE
///
/// Gets the stack trace at the time of allocation for debugging purposes.
///
public StackTrace? StackTrace
{
get; init;
}
#endif
}
public readonly struct AllocationManagerInitOpts
{
public nuint ArenaCapacity
{
get; init;
}
public nuint StackCapacity
{
get; init;
}
public int FreeListConcurrencyLevel
{
get; init;
}
public static AllocationManagerInitOpts Default => new AllocationManagerInitOpts
{
ArenaCapacity = 1024 * 1024 * 1024, // 1 GB
StackCapacity = 16 * 1024 * 1024, // 16 MB per thread
FreeListConcurrencyLevel = Environment.ProcessorCount
};
}
///
/// Provides memory allocation management for native memory allocations, with support for tracking,
/// debugging, and custom allocation strategies.
///
public static unsafe class AllocationManager
{
private struct ArenaAllocator : IAllocator, IDisposable
{
private const int _ARENA_MAGIC_ID = -3941029;
private VirtualArena _arena;
private int _currentTick;
private AllocationHandle _handle;
public readonly AllocationHandle Handle => _handle;
public readonly int CurrentTick => _currentTick;
public void Init(nuint capacity)
{
_arena = new VirtualArena(capacity);
_handle = new AllocationHandle
{
State = Unsafe.AsPointer(ref this),
Alloc = &Allocate,
Realloc = &Reallocate,
Free = null,
#if MHP_ENABLE_SAFETY_CHECKS
IsValid = &IsValid
#else
IsValid = null
#endif
};
_currentTick = 0;
}
private static void* Allocate(void* instance, nuint size, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
var selfPtr = (ArenaAllocator*)instance;
var ptr = selfPtr->_arena.Allocate(size, alignment, allocationOption);
if (ptr == null)
{
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = MemoryHandle.Invalid;
#endif
return null;
}
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = new MemoryHandle(_ARENA_MAGIC_ID, selfPtr->_currentTick);
#endif
return ptr;
}
private static void* Reallocate(void* instance, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
if (ptr == null)
{
return Allocate(instance, newSize, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, pHandle
#endif
);
}
var selfPtr = (ArenaAllocator*)instance;
var newPtr = selfPtr->_arena.Allocate(newSize, alignment, allocationOption);
if (newPtr == null)
{
return null;
}
MemCpy(newPtr, ptr, Math.Min(oldSize, newSize));
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = new MemoryHandle(_ARENA_MAGIC_ID, selfPtr->_currentTick);
#endif
return newPtr;
}
#if MHP_ENABLE_SAFETY_CHECKS
private static bool IsValid(void* instance, MemoryHandle handle)
{
var selfPtr = (ArenaAllocator*)instance;
return handle.ID == _ARENA_MAGIC_ID && handle.Generation == selfPtr->_currentTick;
}
#endif
public void Reset()
{
_arena.Reset();
_currentTick++;
}
public void Dispose()
{
_arena.Dispose();
}
}
private struct HeapAllocator : IAllocator
{
private AllocationHandle _handle;
public readonly AllocationHandle Handle => _handle;
public void Init()
{
_handle = new AllocationHandle
{
State = null,
Alloc = &Allocate,
Realloc = &Reallocate,
Free = &Free,
#if MHP_ENABLE_SAFETY_CHECKS
IsValid = &IsValid
#else
IsValid = null
#endif
};
}
private static void* Allocate(void* _, nuint size, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
return HeapAlloc(size, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, pHandle
#endif
);
}
private static void* Reallocate(void* _, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
if (ptr == null)
{
return Allocate(null, newSize, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, pHandle
#endif
);
}
#if MHP_ENABLE_SAFETY_CHECKS
MemoryHandle newHandle;
#endif
var newPtr = HeapAlloc(newSize, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, &newHandle
#endif
);
if (newPtr == null)
{
return null;
}
MemCpy(newPtr, ptr, Math.Min(oldSize, newSize));
HeapFree(ptr
#if MHP_ENABLE_SAFETY_CHECKS
, *pHandle
#endif
);
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = newHandle;
#endif
return newPtr;
}
private static void Free(void* _, void* ptr
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle handle
#endif
)
{
HeapFree(ptr
#if MHP_ENABLE_SAFETY_CHECKS
, handle
#endif
);
}
#if MHP_ENABLE_SAFETY_CHECKS
private static bool IsValid(void* _, MemoryHandle handle)
{
return ContainsAllocation(handle);
}
#endif
}
private struct StackAllocator : IAllocator, IDisposable
{
private const int _STACK_MAGIC_ID = -6843541;
private static void** s_pStackBuffers = null;
private static int s_stackCount = 0;
private static int s_stackCapacity = 0;
private static readonly SpinLock s_locker = new SpinLock(false);
[ThreadStatic]
private static VirtualStack s_stack;
private AllocationHandle _handle;
public readonly AllocationHandle Handle => _handle;
public void Init()
{
_handle = new AllocationHandle
{
State = Unsafe.AsPointer(ref this),
Alloc = &Allocate,
Realloc = &Reallocate,
Free = null,
#if MHP_ENABLE_SAFETY_CHECKS
IsValid = &IsValid
#else
IsValid = null
#endif
};
}
private static void EnsureInitialize()
{
if (s_stack.Buffer == null)
{
s_stack = new VirtualStack(s_threadLocalStackDefaultSize);
var token = false;
try
{
s_locker.Enter(ref token);
if (s_pStackBuffers == null)
{
s_pStackBuffers = (void**)Malloc((nuint)(sizeof(void*) * Environment.ProcessorCount));
s_stackCapacity = Environment.ProcessorCount;
}
if (s_stackCount >= s_stackCapacity)
{
var pOld = s_pStackBuffers;
var newCapacity = s_stackCapacity * 2;
var pNew = (void**)Realloc(pOld, (nuint)(sizeof(void*) * newCapacity));
s_pStackBuffers = pNew;
s_stackCapacity = newCapacity;
}
s_pStackBuffers[s_stackCount] = s_stack.Buffer;
s_stackCount++;
}
finally
{
if (token)
{
s_locker.Exit();
}
}
}
}
private static void* Allocate(void* instance, nuint size, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
EnsureInitialize();
var ptr = s_stack.Allocate(size, alignment, allocationOption);
if (ptr == null)
{
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = MemoryHandle.Invalid;
#endif
return null;
}
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = new MemoryHandle(_STACK_MAGIC_ID, (int)s_stack.Offset);
#endif
return ptr;
}
private static void* Reallocate(void* instance, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
if (ptr == null)
{
return Allocate(instance, newSize, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, pHandle
#endif
);
}
EnsureInitialize();
// Optimize for last allocation. Set offset directly.
var oldBase = s_stack.Buffer + s_stack.Offset - oldSize;
if (ptr == oldBase)
{
if (newSize > oldSize)
{
var diff = newSize - oldSize;
s_stack.Offset += diff;
if (allocationOption.HasFlag(AllocationOption.Clear))
{
MemClear(s_stack.Buffer + s_stack.Offset - diff, diff);
}
}
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = new MemoryHandle(_STACK_MAGIC_ID, (int)s_stack.Offset);
#endif
return ptr;
}
var newPtr = s_stack.Allocate(newSize, alignment, allocationOption);
if (newPtr == null)
{
return null;
}
MemCpy(newPtr, ptr, Math.Min(oldSize, newSize));
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = new MemoryHandle(_STACK_MAGIC_ID, (int)s_stack.Offset);
#endif
return newPtr;
}
#if MHP_ENABLE_SAFETY_CHECKS
private static bool IsValid(void* instance, MemoryHandle handle)
{
return handle.ID == _STACK_MAGIC_ID && handle.Generation <= (int)s_stack.Offset;
}
#endif
public static VirtualStack.Scope CreateScope(StackAllocator* pSelf)
{
EnsureInitialize();
return s_stack.CreateScope(pSelf->_handle);
}
public readonly void Dispose()
{
if (s_pStackBuffers == null)
{
return;
}
for (var i = 0; i < s_stackCount; i++)
{
Munmap(s_pStackBuffers[i], s_threadLocalStackDefaultSize);
}
}
}
private struct FreeListAllocator : IAllocator, IDisposable
{
private FreeList _freeList;
private AllocationHandle _handle;
public readonly AllocationHandle Handle => _handle;
public void Init(int concurrencyLevel)
{
_freeList = new FreeList(8, 64 * 1024, concurrencyLevel);
_handle = new AllocationHandle
{
State = Unsafe.AsPointer(ref this),
Alloc = &Allocate,
Realloc = &Reallocate,
Free = &Free,
#if MHP_ENABLE_SAFETY_CHECKS
IsValid = &IsValid
#else
IsValid = null
#endif
};
}
private static void* Allocate(void* instance, nuint size, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
var selfPtr = (FreeListAllocator*)instance;
var ptr = selfPtr->_freeList.Allocate(size, alignment, allocationOption);
if (ptr == null)
{
return null;
}
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = AddAllocation(ptr, size);
#endif
return ptr;
}
private static void* Reallocate(void* instance, void* ptr, nuint oldSize, nuint newSize, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
if (ptr == null)
{
return Allocate(instance, newSize, alignment, allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, pHandle
#endif
);
}
var selfPtr = (FreeListAllocator*)instance;
var newPtr = selfPtr->_freeList.Allocate(newSize, alignment, allocationOption);
if (newPtr == null)
{
return null;
}
MemCpy(newPtr, ptr, Math.Min(oldSize, newSize));
selfPtr->_freeList.Free(ptr);
#if MHP_ENABLE_SAFETY_CHECKS
RemoveAllocation(*pHandle);
*pHandle = AddAllocation(newPtr, newSize);
#endif
return newPtr;
}
#if MHP_ENABLE_SAFETY_CHECKS
private static bool IsValid(void* instance, MemoryHandle handle)
{
return ContainsAllocation(handle);
}
#endif
private static void Free(void* instance, void* ptr
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle handle
#endif
)
{
var selfPtr = (FreeListAllocator*)instance;
selfPtr->_freeList.Free(ptr);
#if MHP_ENABLE_SAFETY_CHECKS
RemoveAllocation(handle);
#endif
}
public void Dispose()
{
_freeList.Dispose();
}
}
private static ArenaAllocator* s_pArenaAllocator;
private static HeapAllocator* s_pHeapAllocator;
private static StackAllocator* s_pStackAllocator;
private static FreeListAllocator* s_pFreeListAllocator;
#if MHP_ENABLE_SAFETY_CHECKS
private static ConcurrentSlotMap s_allocations = null!;
#endif
///
/// Gets the number of live tracked heap allocations. Always returns 0 if MHP_ENABLE_SAFETY_CHECKS is disabled.
///
public static int LiveAllocationCount =>
#if MHP_ENABLE_SAFETY_CHECKS
s_allocations.Count;
#else
0;
#endif
private static volatile bool s_initialized;
private static nuint s_threadLocalStackDefaultSize;
public static void Initialize(AllocationManagerInitOpts opts)
{
if (s_initialized)
{
return;
}
#if MHP_ENABLE_SAFETY_CHECKS
s_allocations = new ConcurrentSlotMap(256);
#endif
var ptr = (byte*)Malloc((nuint)(sizeof(ArenaAllocator) + sizeof(HeapAllocator) + sizeof(StackAllocator) + sizeof(FreeListAllocator)));
s_pArenaAllocator = (ArenaAllocator*)ptr;
s_pHeapAllocator = (HeapAllocator*)(ptr + sizeof(ArenaAllocator));
s_pStackAllocator = (StackAllocator*)(ptr + sizeof(ArenaAllocator) + sizeof(HeapAllocator));
s_pFreeListAllocator = (FreeListAllocator*)(ptr + sizeof(ArenaAllocator) + sizeof(HeapAllocator) + sizeof(StackAllocator));
s_pArenaAllocator->Init(opts.ArenaCapacity);
s_pHeapAllocator->Init();
s_pStackAllocator->Init();
s_pFreeListAllocator->Init(opts.FreeListConcurrencyLevel);
s_threadLocalStackDefaultSize = opts.StackCapacity;
s_initialized = true;
}
///
/// Gets a reference to the allocation pHandle for the specified allocator type.
///
/// The allocator type for which to retrieve the allocation pHandle.
/// A reference to the allocation pHandle associated with the specified allocator type.
///
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static AllocationHandle GetAllocationHandle(Allocator allocator)
{
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
return allocator switch
{
Allocator.Temp => s_pArenaAllocator->Handle,
Allocator.Persistent => s_pHeapAllocator->Handle,
Allocator.FreeList => s_pFreeListAllocator->Handle,
_ => throw new ArgumentException("Target allocator type does not support custom allocation.", nameof(allocator)),
};
}
///
/// Allocates a block of memory from the heap with the specified newSize and alignment, using the given allocation options.
///
/// The number of bytes to allocate. Must be greater than zero.
/// The alignment, in bytes, for the allocated memory block. Must be a power of two.
/// An optional set of flags that control allocation behavior, such as whether the memory should be cleared or
/// tracked. The default is .
/// A pointer to the beginning of the allocated memory block.
/// Thrown if the allocation fails.
public static void* HeapAlloc(nuint size, nuint alignment, AllocationOption allocationOption
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle* pHandle
#endif
)
{
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
var ptr = AlignedAlloc(size, alignment);
if (ptr == null)
{
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = MemoryHandle.Invalid;
#endif
return null;
}
if (allocationOption.HasFlag(AllocationOption.Clear))
{
MemClear(ptr, size);
}
#if MHP_ENABLE_SAFETY_CHECKS
*pHandle = AddAllocation(ptr, size);
#endif
return ptr;
}
///
/// Releases a block of unmanaged memory previously allocated by the heap allocator.
///
/// A pointer to the memory block to be freed. The pointer must have been returned by a compatible heap allocation
/// method and must not be null.
/// The handle representing the memory allocation to free. The handle must be valid and previously allocated.
public static void HeapFree(void* ptr
#if MHP_ENABLE_SAFETY_CHECKS
, MemoryHandle handle
#endif
)
{
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
AlignedFree(ptr);
#if MHP_ENABLE_SAFETY_CHECKS
RemoveAllocation(handle);
#endif
}
///
/// Releases a block of unmanaged memory previously allocated by the heap allocator.
///
///
/// No ops when MHP_ENABLE_SAFETY_CHECKS is disabled, as we cannot fetch the allocation info from the handle to get the pointer to free.
///
/// The handle representing the memory allocation to free. The handle must be valid and previously allocated.
public static void HeapFree(MemoryHandle handle)
{
#if MHP_ENABLE_SAFETY_CHECKS
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
if (TryGetAllocation(handle, out var info))
{
HeapFree((void*)info.Address, handle);
}
#endif
// No-op when safety checks are disabled, as we cannot fetch the allocation info from the handle.
}
///
/// Resets the temporary memory allocator, clearing all allocated memory.
///
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static void ResetTempAllocator()
{
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
s_pArenaAllocator->Reset();
}
///
/// Creates a new thread local stack scope for managing temporary allocations within the current context.
///
/// A instance representing the newly created stack scope. The scope must be disposed when no longer needed to release allocated resources.
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static VirtualStack.Scope CreateStackScope()
{
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
return StackAllocator.CreateScope(s_pStackAllocator);
}
///
/// Registers a memory allocation and returns a handle that can be used to manage or reference the allocated memory.
///
///
/// Always returns an invalid handle if MHP_ENABLE_SAFETY_CHECKS is disabled.
///
/// A pointer to the memory block to be registered. The pointer must reference a valid, allocated memory region.
/// The newSize of the memory block to be registered.
/// A MemoryHandle representing the registered allocation.
[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
};
var id = s_allocations.Add(info, out var generation);
return new MemoryHandle(id, generation);
#else
return MemoryHandle.Invalid;
#endif
}
///
/// Removes the memory allocation associated with the specified handle.
///
///
/// Always returns false if debug layer is disabled.
///
/// The handle representing the memory allocation to remove. The handle must be valid and previously allocated.
/// true if the allocation was successfully removed; otherwise, false.
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static bool RemoveAllocation(MemoryHandle handle)
{
#if MHP_ENABLE_SAFETY_CHECKS
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
return s_allocations.Remove(handle.ID, handle.Generation, out var info);
#else
return false;
#endif
}
///
/// Attempts to retrieve the memory allocation pointer associated with the specified handle.
///
///
/// Always returns false if debug layer is disabled, and the output pointer will be set to .
///
/// The memory handle identifying the allocation to retrieve allocation.
/// When this method returns, contains the allocation information associated with the specified handle, if the allocation was found; otherwise, an uninitialized value.
/// true if the allocation was found and contains valid allocation information; otherwise, false.
[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
}
///
/// Determines whether the specified memory handle refers to a currently tracked allocation.
///
///
/// This only validates the memory when you added the allocation via .
/// For validating memory from , use instead.
/// Always returns false if debug layer is disabled.
///
/// The memory handle to check for an associated allocation.
/// true if the allocation corresponding to the handle exists; otherwise, false.
[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
}
///
/// Gets the total newSize of all currently tracked allocations.
///
///
/// Always returns 0 if MHP_ENABLE_SAFETY_CHECKS is disabled.
///
/// The total newSize of all currently tracked allocations.
public static nuint GetTotalAllocatedMemory()
{
#if MHP_ENABLE_SAFETY_CHECKS
Debug.Assert(s_initialized, "AllocationManager is not initialized.");
nuint total = 0;
foreach (var allocation in s_allocations)
{
total += allocation.Size;
}
return total;
#else
return 0;
#endif
}
///
/// Disposes of the AllocationManager, freeing all allocated memory and resources.
///
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_pArenaAllocator->Dispose();
s_pStackAllocator->Dispose();
s_pFreeListAllocator->Dispose();
Free(s_pArenaAllocator);
}
}