WIP: Removed internal header tracking for WorkspaceHeap. Modified Handles to track the alignment that the underlying memory was aligned to. Re-worked WorkspaceHeap to use lists rather than Stack

UnmangedMMU/Handles/Internal/IOwnedHandle.cs

@@ -11,5 +11,10 @@ namespace UnmanagedMMU.Handles.Internal
         /// Returns the <see cref="IUnmanagedMemoryOwner"/> that created owns this <see cref="IOwnedHandle"/>
         /// </summary>
         IUnmanagedMemoryOwner GetOwner();
+
+        /// <summary>
+        /// Returns the memory alignment for the <see cref="IOwnedHandle"/>
+        /// </summary>
+        nuint Alignment { get; }
     }
 }

UnmangedMMU/Handles/MemoryHandleBase.cs

@@ -41,18 +41,25 @@ namespace UnmanagedMMU.Handles
         /// </summary>
         private bool _disposed;
 
+        /// <summary>
+        /// The alignment of the memory pointed to by <see cref="_ptr"/>
+        /// </summary>
+        private readonly nuint _alignment;
+
         /// <summary>
         /// Initializes a new <see cref="MemoryHandleBase{T}"/> instnace
         /// </summary>
         /// <param name="ptr">Pointer to the allocated unmanaged memory</param>
         /// <param name="byteLength">The size of the unallocated memory block in bytes</param>
+        /// <param name="alignment">The alignment that the unmanged memory pointed to by <paramref name="ptr"/> </param>
         /// <param name="owner">The <see cref="IUnmanagedMemoryOwner"/> that owns the <see cref="MemoryHandleBase{T}"/> handle being created</param>
-        protected MemoryHandleBase(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner)
+        protected MemoryHandleBase(void* ptr, nuint byteLength, nuint alignment, IUnmanagedMemoryOwner owner)
         {
             // Defensive check
             Debug.Assert(ptr != null, message: "BUG CHECK: E_INVALID_MEMORY_HANDLE");
             _ptr = ptr;
             _bytelen = byteLength;
+            _alignment = alignment;
             _owner = owner;
         }
 
@@ -64,6 +71,14 @@ namespace UnmanagedMMU.Handles
             get { return _ptr; }
         }
 
+        /// <summary>
+        /// Returns the memory alignment for the <see cref="MemoryHandleBase{T}"/>
+        /// </summary>
+        public virtual nuint Alignment
+        {
+            get { return _alignment; }
+        }
+
         /// <summary>
         /// Gets the typed pointer to the unmanged memory block
         /// </summary>

UnmangedMMU/Handles/PersistentMemoryHandle.cs

@@ -5,7 +5,7 @@ namespace UnmanagedMMU.Handles
     internal sealed unsafe class PersistentMemoryHandle<T> : MemoryHandleBase<T> where T : unmanaged
     {
 
-        public PersistentMemoryHandle(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, owner)
+        public PersistentMemoryHandle(void* ptr, nuint byteLength, nuint alignment, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, alignment, owner)
         {
         }
 

UnmangedMMU/Handles/SegmentedMemoryHandle.cs

@@ -1,15 +1,10 @@
-using System;
+using UnmanagedMMU.Allocators;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
-using UnmanagedMMU.Allocators;
 
 namespace UnmanagedMMU.Handles
 {
     internal sealed unsafe class SegmentedMemoryHandle<T> : MemoryHandleBase<T> where T : unmanaged
     {
-        public SegmentedMemoryHandle(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, owner)
+        public SegmentedMemoryHandle(void* ptr, nuint byteLength, nuint alignment, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, alignment, owner)
         {
         }
 

UnmangedMMU/SegmentedPool.cs

@@ -291,7 +291,7 @@
         /// <exception cref="ArgumentException">
         /// Thrown if <paramref name="segmentSize"/> is zero, or if <paramref name="initialSegments"/> is less than 1.
         /// </exception>
-        public SegmentedPool(nuint segmentSize = _defaultSegmentSize, SegmentAlignment segmentAlignment = SegmentAlignment.Aligned32, int initialSegments = 4, bool zeroMemory = false)
+        public SegmentedPool(nuint segmentSize = _defaultSegmentSize, SegmentAlignment segmentAlignment = SegmentAlignment.Aligned16, int initialSegments = 4, bool zeroMemory = false)
             : this(segmentSize, segmentAlignment, initialSegments, zeroMemory, new DefaultUnmanagedAllocator())
         {
         }
@@ -620,7 +620,9 @@
         {
             T* ptr = Alloc<T>(count);
             nuint byteLength = (nuint)count * (nuint)sizeof(T);
-            return new SegmentedMemoryHandle<T>(ptr, byteLength, this);
+
+            nuint alignment = _segmentAlignment > (nuint)sizeof(T) ? _segmentAlignment : (nuint)sizeof(T);
+            return new SegmentedMemoryHandle<T>(ptr, byteLength, alignment, this);
         }
 
         /// <summary>
@@ -651,7 +653,7 @@
 
             T* ptr = AllocateWithAlignment<T>(count, effectiveAlignment);
             nuint byteLength = (nuint)count * (nuint)sizeof(T);
-            return new SegmentedMemoryHandle<T>(ptr, byteLength, this);
+            return new SegmentedMemoryHandle<T>(ptr, byteLength, effectiveAlignment, this);
 
         }
 

UnmangedMMU/WorkspaceHeap.cs

@@ -3,7 +3,6 @@
     using System;
     using System.Collections.Generic;
     using System.Runtime.CompilerServices;
-    using System.Runtime.InteropServices;
     using UnmanagedMMU.Allocators;
     using UnmanagedMMU.Handles;
     using UnmanagedMMU.Handles.Internal;
@@ -33,6 +32,16 @@
     /// </remarks>
     public unsafe sealed class WorkspaceHeap : IDisposable, IUnmanagedMemoryOwner
     {
+
+        /// <summary>
+        /// Struct defining metadata to be stored with each allocation
+        /// </summary>
+        private struct AllocationMetadata
+        {
+            public IntPtr Block;
+            public nuint Alignment;
+        }
+
         /// <summary>
         /// The maximum size, in bytes, for "small" allocations. Uses fixed-size buckets
         /// </summary>
@@ -61,12 +70,12 @@
         /// Predefined bucket sizes used for small allocation reuse.
         /// </summary>
         private static readonly nuint[] _sizeClasses =
-        {
+        [
              32,  64,  96, 128, 160, 192, 224, 256,
             288, 320, 352, 384, 416, 448, 480, 512,
             544, 576, 608, 640, 672, 704, 736, 768,
             800, 832, 864, 896, 928, 960, 992, 1024
-        };
+        ];
 
         /// <summary>
         /// Allocator interface used for all underlying unmanaged memory operations.
@@ -84,17 +93,17 @@
         /// <remarks>
         /// For a given bucket, the corresponding stack contains previously allocated blocks that are available to be used
         /// </remarks>
-        private readonly Dictionary<nuint, Stack<IntPtr>> _smallFree = new();
+        private readonly Dictionary<nuint, List<AllocationMetadata>> _smallFree = [];
 
         /// <summary>
         /// Free lists for medium allocations keyed by the exactly allocated size, in bytes, and sorted for best-fit.
         /// </summary>
-        private readonly SortedDictionary<nuint, Stack<IntPtr>> _mediumFree = new();
+        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _mediumFree = [];
 
         /// <summary>
         /// Free lists for large allocations keyed by the exactly allocated size, in bytes, sorted for tolerance-based reuse.
         /// </summary>
-        private readonly SortedDictionary<nuint, Stack<IntPtr>> _largeFree = new();
+        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _largeFree = [];
 
         /// <summary>
         /// Tracks the total bytes of memory reserved from the provided <see cref="IUnmanagedAllocator"/>.
@@ -116,33 +125,6 @@
         /// </summary>
         private volatile bool _disposed;
 
-        /// <summary>
-        /// Internal header prepended to each allocation to track its size.
-        /// </summary>
-        [StructLayout(LayoutKind.Sequential)]
-        private struct BlockHeader
-        {
-            /// <summary>
-            /// Size of the allocation in bytes.
-            /// </summary>
-            public nuint Size;
-
-            /// <summary>
-            /// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
-            /// </summary>
-            private readonly nuint _pad1;
-
-            /// <summary>
-            /// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
-            /// </summary>
-            private readonly nuint _pad2;
-
-            /// <summary>
-            /// Padding to ensure <see cref="BlockHeader"></see> is 32-byte aligned
-            /// </summary>
-            private readonly nuint _pad3;
-        }
-
         /// <summary>
         /// Creates a new <see cref="WorkspaceHeap"/>.
         /// </summary>
@@ -160,8 +142,10 @@
             _allocator = allocator;
 
             // Initialize small-size buckets
-            foreach (var size in _sizeClasses)
+            foreach (nuint size in _sizeClasses)
-                _smallFree[size] = new Stack<IntPtr>();
+            {
+                _smallFree[size] = [];                    
+            }
         }
 
         /// <summary>
@@ -243,16 +227,15 @@
         /// Allocates a new block from the underlying allocator including a header.
         /// </summary>
         /// <param name="payloadSize">Requested payload size in bytes.</param>
+        /// <param name="alignment">FOO</param>
         /// <returns>
         /// Pointer to the allocated block (header included).
         /// </returns>
-        private IntPtr AllocateNew(nuint payloadSize)
+        private IntPtr AllocateNewAligned(nuint payloadSize, nuint alignment)
         {
-            nuint total = payloadSize + (nuint)sizeof(BlockHeader);
+            void* raw = _allocator.AllocAligned(payloadSize, alignment);
-            void* raw = _allocator.Alloc(total);
+            _totalReserved += payloadSize;
-            _totalReserved += total;
             _totalAllocations++;
-
             return (IntPtr)raw;
         }
 
@@ -264,100 +247,154 @@
         /// <returns> An <see cref="IMemoryHandle{T}"/> to the allocated memory.</returns>
         /// <exception cref="ObjectDisposedException">Thrown if heap is disposed.</exception>
         /// <exception cref="ArgumentOutOfRangeException">Thrown if size is zero.</exception>
+        /// <remarks>The memory returned by this method is always 16-byte aligned, for other alignemnts use </remarks>
         public IMemoryHandle<T> Allocate<T>(int count, bool zero = false) where T : unmanaged
+        {
+            return AllocateAligned<T>(count, SegmentAlignment.Aligned16, zero);
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="count"></param>
+        /// <param name="alignment"></param>
+        /// <param name="zero"></param>
+        /// <returns></returns>
+        /// <exception cref="OverflowException"></exception>
+        public IMemoryHandle<T> AllocateAligned<T>(int count, SegmentAlignment alignment, bool zero = false) where T: unmanaged
         {
             ArgumentOutOfRangeException.ThrowIfNegative(count);
             ArgumentOutOfRangeException.ThrowIfZero(count);
-
             if ((nuint)count > nuint.MaxValue / (nuint)(sizeof(T)))
             {
                 throw new OverflowException($"Requested allocation of {count} elements of type {typeof(T)} exceeds allowable maximum memory size.");
             }
 
             nuint size = (nuint)count * (nuint)sizeof(T);
+            nuint requestedAlignment = (nuint)alignment;
+            nuint effectiveAlignment = requestedAlignment < (nuint)sizeof(T) ? (nuint)sizeof(T) : requestedAlignment;
 
             lock (_lock)
             {
                 ThrowIfDisposed();
-                void* ptr = null;
+                void* ptr;
-                if (size <= _smallThreshold)
-                {
-                    ptr = AllocateSmall(size, zero);
-                }
-                else if (size <= _mediumThreshold)
-                {
-                    ptr = AllocateMedium(size, zero);
-                }
-                else
-                {
-                    ptr = AllocateLarge(size, zero);
-                }
 
-                return new PersistentMemoryHandle<T>((T*)ptr, size, this);
+                if (size <= _smallThreshold)
+                    ptr = AllocateSmallAligned(size, effectiveAlignment, zero);
+                else if (size <= _mediumThreshold)
+                    ptr = AllocateMediumAligned(size, effectiveAlignment, zero);
+                else
+                    ptr = AllocateLargeAligned(size, effectiveAlignment, zero);
+
+                return new PersistentMemoryHandle<T>((T*)ptr, size, effectiveAlignment, this);
             }
         }
 
-        /// <summary>Allocates a small-size block using bucketed free lists.</summary>
+        /// <summary>
-        private void* AllocateSmall(nuint size, bool zero)
+        /// TODO: Fill in
+        /// </summary>
+        /// <param name="alignment"></param>
+        /// <param name="list"></param>
+        /// <returns></returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static int FindIndexAlignmentMatch(nuint alignment, List<AllocationMetadata> list)
+        {
+            // Attempt to find an index match the requested size
+            int mIdx = -1;
+            for (int i = 0; i < list.Count; i++)
+            {
+                if (list[i].Alignment >= alignment)
+                {
+                    mIdx = i;
+                    break;
+                }
+            }
+            return mIdx;
+        }
+
+        /// <summary>
+        /// Allocates a small-size block using bucketed free lists.
+        /// </summary>
+        private void* AllocateSmallAligned(nuint size, nuint alignment, bool zero)
         {
             nuint bucket = GetSizeClass(size);
-            Stack<IntPtr> stack = _smallFree[bucket];
+            List<AllocationMetadata> free = _smallFree[bucket];
 
-            IntPtr block = stack.Count > 0
+            // Attempt to find an index match the requested size
-                ? stack.Pop()
+            int mIdx = FindIndexAlignmentMatch(alignment, free);
-                : AllocateNew(bucket);
 
-            BlockHeader* header = (BlockHeader*)block;
+            if (mIdx >= 0)
-            header->Size = bucket;
+            {
+                AllocationMetadata meta = free[mIdx];
+                free.RemoveAt(mIdx);
 
+                _totalInUse += bucket;
+                void* user = (void*)meta.Block;
+                if (zero)
+                {
+                    Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
+                }
+                return user;
+            }
+
+            // allocate a new small block
+            IntPtr block = AllocateNewAligned(bucket, alignment);
             _totalInUse += bucket;
-
+            void* nUser = (void*)block;
-            void* user = header + 1;
-
             if (zero)
-                Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
+            {
-
+                Unsafe.InitBlockUnaligned(nUser, 0, (uint)bucket);
-            return user;
+            }
+            return nUser;
         }
 
         /// <summary>Allocates a medium-size block using best-fit reuse.</summary>
-        private void* AllocateMedium(nuint size, bool zero)
+        private void* AllocateMediumAligned(nuint size, nuint alignment, bool zero)
         {
-            foreach (var kv in _mediumFree)
+            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _mediumFree)
             {
                 if (kv.Key >= size && kv.Value.Count > 0)
                 {
-                    var block = kv.Value.Pop();
+                    List<AllocationMetadata> free = kv.Value;
-                    var header = (BlockHeader*)block;
-                    _totalInUse += header->Size;
 
-                    void* user = header + 1;
+                    // Attempt to find an index match the requested size
-                    if (zero)
+                    int mIdx = FindIndexAlignmentMatch(alignment, free);
-                        Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
 
-                    return user;
+                    if (mIdx >= 0)
+                    {
+                        AllocationMetadata meta = free[mIdx];
+                        free.RemoveAt(mIdx);
+
+                        _totalInUse += kv.Key;
+                        void* user = (void*)meta.Block;
+                        if (zero)
+                        {
+                            Unsafe.InitBlockUnaligned(user, 0, (uint)kv.Key);
+                        }
+                        return user;
+                    }
                 }
             }
-
+            // no match
-            var newBlock = AllocateNew(size);
+            IntPtr newBlock = AllocateNewAligned(size, alignment);
-            var newHeader = (BlockHeader*)newBlock;
-            newHeader->Size = size;
             _totalInUse += size;
 
-            void* newUser = newHeader + 1;
+            void* nUser = (void*)newBlock;
             if (zero)
             {
-                Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
+                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
             }               
 
-
+            return nUser;
-            return newUser;
         }
 
-        /// <summary>Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).</summary>
+        /// <summary>
-        private void* AllocateLarge(nuint size, bool zero)
+        /// Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).
+        /// </summary>
+        private void* AllocateLargeAligned(nuint size, nuint alignment, bool zero)
         {
-            foreach (var kv in _largeFree)
+            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _largeFree)
             {
                 nuint blockSize = kv.Key;
                 if (blockSize < size || kv.Value.Count == 0)
@@ -365,41 +402,46 @@
                     continue;
                 }
 
+                // check waste tolerance before attempting to find an alignment match
                 nuint waste = blockSize - size;
-                bool acceptable =
+                bool acceptable = waste <= _largeMaxWasteBytes || ((double)blockSize / size) <= _largeWasteRatioLimit;
-                    waste <= _largeMaxWasteBytes ||
-                    ((double)blockSize / size) <= _largeWasteRatioLimit;
 
                 if (!acceptable)
                 {
                     continue;
                 }
 
-                var block = kv.Value.Pop();
+                List<AllocationMetadata> free = kv.Value;
-                var header = (BlockHeader*)block;
-                _totalInUse += header->Size;
 
-                void* user = header + 1;
+                // Attempt to find an index match the requested size
-                if (zero)
+                int mIdx = FindIndexAlignmentMatch(alignment, free);
+
+                if (mIdx >= 0)
                 {
-                    Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
+                    AllocationMetadata meta = free[mIdx];
-                }
+                    free.RemoveAt(mIdx);
 
-                return user;
+                    _totalInUse += blockSize;
+                    void* user = (void*)meta.Block;
+                    if (zero)
+                    {
+                        Unsafe.InitBlockUnaligned(user, 0, (uint)blockSize);
+                    }
+
+                    return user;
+                }
             }
 
-            var newBlock = AllocateNew(size);
+            // no match
-            var newHeader = (BlockHeader*)newBlock;
-            newHeader->Size = size;
-            _totalInUse += size;
 
-            void* newUser = newHeader + 1;
+            IntPtr block = AllocateNewAligned(size, alignment);
+            _totalInUse += size;
+            void* nUser = (void*)block;
             if (zero)
             {
-                Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
+                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
             }
-
+            return nUser;
-            return newUser;
         }
 
         /// <summary>
@@ -421,23 +463,41 @@
 
             lock (_lock)
             {
-                var header = ((BlockHeader*)handle.Pointer) - 1;
+                nuint size = handle.ByteCount;
-                nuint size = header->Size;
+                nuint alignment = handle.Alignment;
                 _totalInUse -= size;
 
                 if (size <= _smallThreshold)
-                    _smallFree[size].Push((IntPtr)header);
+                {
+                    _smallFree[size].Add(new AllocationMetadata
+                    {
+                        Block = (IntPtr)handle.Pointer,
+                        Alignment = alignment
+                    });
+                }
                 else if (size <= _mediumThreshold)
                 {
-                    if (!_mediumFree.TryGetValue(size, out var stack))
+                    if (!_mediumFree.TryGetValue(size, out List<AllocationMetadata>? free))
-                        _mediumFree[size] = stack = new Stack<IntPtr>();
+                    {
-                    stack.Push((IntPtr)header);
+                        _mediumFree[size] = free = [];
+                    }
+                    free.Add(new AllocationMetadata
+                    {
+                        Block = (IntPtr)handle.Pointer,
+                        Alignment = alignment
+                    });
                 }
                 else
                 {
-                    if (!_largeFree.TryGetValue(size, out var stack))
+                    if (!_largeFree.TryGetValue(size, out List<AllocationMetadata>? free))
-                        _largeFree[size] = stack = new Stack<IntPtr>();
+                    {
-                    stack.Push((IntPtr)header);
+                        _largeFree[size] = free = [];
+                    }
+                    free.Add(new AllocationMetadata
+                    {
+                        Block = (IntPtr)handle.Pointer,
+                        Alignment = alignment
+                    });
                 }
             }
         }
@@ -457,17 +517,17 @@
         }
 
         /// <summary>Helper to free all blocks in a dictionary of free stacks.</summary>
-        private void PruneDictionary(IDictionary<nuint, Stack<IntPtr>> dict)
+        private void PruneDictionary(IDictionary<nuint, List<AllocationMetadata>> dict)
         {
-            foreach (var kv in dict)
+            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in dict)
             {
-                var stack = kv.Value;
+                List<AllocationMetadata> list = kv.Value;
-                while (stack.Count > 0)
+                while (list.Count > 0)
                 {
-                    var block = stack.Pop();
+                    var item = list[^1];
-                    var header = (BlockHeader*)block;
+                    list.RemoveAt(list.Count - 1);
-                    _allocator.Free((void*)block);
+                    _allocator.FreeAligned((void*)item.Block, item.Alignment);
-                    _totalReserved -= (header->Size + (nuint)sizeof(BlockHeader));
+                    _totalReserved -= kv.Key;
                 }
             }
         }