Refactor memory management and allocation strategies

Changed the `AllocationManager` initialization to specify size in `Program.cs` and commented out the disposal of `unfreeArray`.
Changed `UnsafeArray` to use `AllocationManager.Realloc` for memory reallocation.
Changed constructors of `UnsafeHashMap` and `UnsafeHashSet` to include an `AllocationOption` parameter for improved memory management.
Added a `Remove` method to `UnsafeHashMap` for key-value pair removal.
Changed `AllocationManager` to use a hash map instead of a queue for managing allocations.
Added a `Realloc` method in `AllocationManager` for handling memory reallocation.
Changed the `Free` method in `AllocationManager` to iterate through the hash map for freeing allocated memory.

Misaki.HighPerformance.Test/Program.cs

@@ -1,7 +1,8 @@
 using Misaki.HighPerformance.Unsafe.Collections;
 using Misaki.HighPerformance.Unsafe.Services;
 
+AllocationManager.Initialize(100);
 var unfreeArray = new UnsafeArray<int>(10, Allocator.Persistent);
-unfreeArray.Dispose();
+//unfreeArray.Dispose();
 
 AllocationManager.Dispose();

Misaki.HighPerformance.Unsafe/Collections/UnsafeArray.cs

@@ -128,7 +128,7 @@ public unsafe struct UnsafeArray<T> : IUnsafeCollection<T>
             return;
         }
 
-        _buffer = (T*)AlignedRealloc(_buffer, (nuint)(newSize * sizeof(T)), AlignOf<T>());
+        _buffer = AllocationManager.Realloc<T>(_buffer, (uint)newSize, (uint)AlignOf<T>(), _allocator);
         _count = newSize;
     }
 

Misaki.HighPerformance.Unsafe/Collections/UnsafeHashMap.cs

@@ -92,9 +92,9 @@ public unsafe struct UnsafeHashMap<TKey, TValue> : IUnsafeCollection<KeyValuePai
     public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator() => new Enumerator((HashMapHelper<TKey>*)UnsafeUtilities.AddressOf(ref _hashMap));
     IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
-    public UnsafeHashMap(int capacity, Allocator allocator)
+    public UnsafeHashMap(int capacity, Allocator allocator, AllocationOption allocationOption = AllocationOption.None)
     {
-        _hashMap = new HashMapHelper<TKey>(capacity, sizeof(TValue), HashMapHelper<TKey>.MINIMAL_CAPACITY, allocator);
+        _hashMap = new HashMapHelper<TKey>(capacity, sizeof(TValue), HashMapHelper<TKey>.MINIMAL_CAPACITY, allocator, allocationOption);
     }
 
     /// <summary>
@@ -132,6 +132,16 @@ public unsafe struct UnsafeHashMap<TKey, TValue> : IUnsafeCollection<KeyValuePai
         }
     }
 
+    /// <summary>
+    /// Removes a particular key and its value.
+    /// </summary>
+    /// <param name="item">The value to remove.</param>
+    /// <returns>True if the value was present.</returns>
+    public bool Remove(TKey key)
+    {
+        return -1 != _hashMap.TryRemove(key);
+    }
+
     /// <summary>
     /// Returns the value associated with a key.
     /// </summary>

Misaki.HighPerformance.Unsafe/Collections/UnsafeHashSet.cs

@@ -49,9 +49,9 @@ public unsafe struct UnsafeHashSet<T> : IUnsafeCollection<T>, IEnumerable<T>
     public IEnumerator<T> GetEnumerator() => new Enumerator((HashMapHelper<T>*)UnsafeUtilities.AddressOf(ref _hashMap));
     IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
-    public UnsafeHashSet(int capacity, Allocator allocator)
+    public UnsafeHashSet(int capacity, Allocator allocator, AllocationOption allocationOption)
     {
-        _hashMap = new HashMapHelper<T>(capacity, 0, HashMapHelper<T>.MINIMAL_CAPACITY, allocator);
+        _hashMap = new HashMapHelper<T>(capacity, 0, HashMapHelper<T>.MINIMAL_CAPACITY, allocator, allocationOption);
     }
 
     /// <summary>

Misaki.HighPerformance.Unsafe/Helpers/HashMapHelper.cs

@@ -128,7 +128,7 @@ public unsafe struct HashMapHelper<TKey> : IDisposable
         return totalSize;
     }
 
-    public HashMapHelper(int capacity, int sizeOfTValue, uint minGrowth, Allocator allocator, AllocationOption allocationOption = AllocationOption.None)
+    public HashMapHelper(int capacity, int sizeOfTValue, uint minGrowth, Allocator allocator, AllocationOption allocationOption)
     {
         if (capacity <= 0)
         {

Misaki.HighPerformance.Unsafe/Services/AllocationManager.cs

@@ -5,18 +5,12 @@ namespace Misaki.HighPerformance.Unsafe.Services;
 
 public static unsafe class AllocationManager
 {
-    private readonly struct AllocationInfo(void* ptr, nuint size)
-    {
-        public readonly void* ptr = ptr;
-        public readonly nuint size = size;
-    }
-
     private const uint _DEFAULT_ARENA_SIZE = 512 * 1024; // 512 KB
 
     private static DynamicArena _arena;
     private static bool _initialized;
 
-    private static UnsafeQueue<AllocationInfo> _allocated;
+    private static UnsafeHashMap<IntPtr, nuint> _allocated;
 
     private static readonly Lock _lock = new();
 
@@ -32,7 +26,7 @@ public static unsafe class AllocationManager
         }
 
         _arena = new DynamicArena(initialSize);
-        _allocated = new UnsafeQueue<AllocationInfo>(32, Allocator.Persistent, AllocationOption.UnTracked);
+        _allocated = new(32, Allocator.Persistent, AllocationOption.UnTracked);
 
         _initialized = true;
     }
@@ -62,7 +56,7 @@ public static unsafe class AllocationManager
                 case Allocator.Persistent:
                     var allocationSize = size * (nuint)sizeof(T);
                     buffer = (T*)AlignedAlloc(allocationSize, alignSize);
-                    _allocated.Enqueue(new AllocationInfo(buffer, allocationSize));
+                    _allocated[(IntPtr)buffer] = allocationSize;
                     break;
 
                 default:
@@ -78,6 +72,42 @@ public static unsafe class AllocationManager
         }
     }
 
+    internal static T* Realloc<T>(T* buffer, uint size, uint alignSize, Allocator allocator)
+        where T : unmanaged
+    {
+        if (!_initialized)
+        {
+            throw new InvalidOperationException("The AllocationManager has not been initialized.");
+        }
+
+        lock (_lock)
+        {
+            T* newBuffer;
+            switch (allocator)
+            {
+                case Allocator.Temp:
+                    newBuffer = (T*)_arena.Allocate(size * (uint)sizeof(T), alignSize, AllocationOption.None);
+                    break;
+
+                case Allocator.Persistent:
+                    var allocationSize = size * (nuint)sizeof(T);
+                    newBuffer = (T*)AlignedRealloc(buffer, allocationSize, alignSize);
+
+                    // If the allocation map can not find the old value, which means that it's a untracked allocation
+                    if (_allocated.Remove((IntPtr)buffer))
+                    {
+                        _allocated.Add((IntPtr)newBuffer, allocationSize);
+                    }
+                    break;
+
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(allocator), "Invalid allocator type.");
+            }
+
+            return newBuffer;
+        }
+    }
+
     internal static void Free(void* ptr, Allocator allocator)
     {
         lock (_lock)
@@ -117,10 +147,10 @@ public static unsafe class AllocationManager
         _arena.Dispose();
 
         nuint unfreeBytes = 0u;
-        while (_allocated.TryDequeue(out var allocationInfo))
+        foreach (var pair in _allocated)
         {
-            unfreeBytes += allocationInfo.size;
-            AlignedFree(allocationInfo.ptr);
+            unfreeBytes += pair.Value;
+            AlignedFree((void*)pair.Key);
         }
 
         _allocated.Dispose();