Updated the SegmentedPool to allow for zeroing of allocations and for returning a pointer to allocated memory to allow for use in unsafe structs

UnmangedMMU/Handles/Internal/IOwnedHandle.cs

@@ -11,10 +11,5 @@ 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,25 +41,18 @@ 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, nuint alignment, IUnmanagedMemoryOwner owner)
+        protected MemoryHandleBase(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner)
         {
             // Defensive check
             Debug.Assert(ptr != null, message: "BUG CHECK: E_INVALID_MEMORY_HANDLE");
             _ptr = ptr;
             _bytelen = byteLength;
-            _alignment = alignment;
             _owner = owner;
         }
 
@@ -71,14 +64,6 @@ 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, nuint alignment, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, alignment, owner)
+        public PersistentMemoryHandle(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, owner)
         {
         }
 

UnmangedMMU/Handles/SegmentedMemoryHandle.cs

@@ -1,10 +1,15 @@
-using UnmanagedMMU.Allocators;
+using System;
+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, nuint alignment, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, alignment, owner)
+        public SegmentedMemoryHandle(void* ptr, nuint byteLength, IUnmanagedMemoryOwner owner) : base(ptr, byteLength, 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.Aligned16, int initialSegments = 4, bool zeroMemory = false)
+        public SegmentedPool(nuint segmentSize = _defaultSegmentSize, SegmentAlignment segmentAlignment = SegmentAlignment.Aligned32, int initialSegments = 4, bool zeroMemory = false)
             : this(segmentSize, segmentAlignment, initialSegments, zeroMemory, new DefaultUnmanagedAllocator())
         {
         }
@@ -620,9 +620,7 @@
         {
             T* ptr = Alloc<T>(count);
             nuint byteLength = (nuint)count * (nuint)sizeof(T);
-
-            nuint alignment = _segmentAlignment > (nuint)sizeof(T) ? _segmentAlignment : (nuint)sizeof(T);
-            return new SegmentedMemoryHandle<T>(ptr, byteLength, alignment, this);
+            return new SegmentedMemoryHandle<T>(ptr, byteLength, this);
         }
 
         /// <summary>
@@ -653,7 +651,7 @@
 
             T* ptr = AllocateWithAlignment<T>(count, effectiveAlignment);
             nuint byteLength = (nuint)count * (nuint)sizeof(T);
-            return new SegmentedMemoryHandle<T>(ptr, byteLength, effectiveAlignment, this);
+            return new SegmentedMemoryHandle<T>(ptr, byteLength, this);
 
         }
 

UnmangedMMU/WorkspaceHeap.cs

@@ -3,6 +3,7 @@
     using System;
     using System.Collections.Generic;
     using System.Runtime.CompilerServices;
+    using System.Runtime.InteropServices;
     using UnmanagedMMU.Allocators;
     using UnmanagedMMU.Handles;
     using UnmanagedMMU.Handles.Internal;
@@ -32,16 +33,6 @@
     /// </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>
@@ -70,12 +61,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.
@@ -93,17 +84,17 @@
         /// <remarks>
         /// For a given bucket, the corresponding stack contains previously allocated blocks that are available to be used
         /// </remarks>
-        private readonly Dictionary<nuint, List<AllocationMetadata>> _smallFree = [];
+        private readonly Dictionary<nuint, Stack<IntPtr>> _smallFree = new();
 
         /// <summary>
         /// Free lists for medium allocations keyed by the exactly allocated size, in bytes, and sorted for best-fit.
         /// </summary>
-        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _mediumFree = [];
+        private readonly SortedDictionary<nuint, Stack<IntPtr>> _mediumFree = new();
 
         /// <summary>
         /// Free lists for large allocations keyed by the exactly allocated size, in bytes, sorted for tolerance-based reuse.
         /// </summary>
-        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _largeFree = [];
+        private readonly SortedDictionary<nuint, Stack<IntPtr>> _largeFree = new();
 
         /// <summary>
         /// Tracks the total bytes of memory reserved from the provided <see cref="IUnmanagedAllocator"/>.
@@ -125,6 +116,33 @@
         /// </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>
@@ -142,10 +160,8 @@
             _allocator = allocator;
 
             // Initialize small-size buckets
-            foreach (nuint size in _sizeClasses)
-            {
-                _smallFree[size] = [];                    
-            }
+            foreach (var size in _sizeClasses)
+                _smallFree[size] = new Stack<IntPtr>();
         }
 
         /// <summary>
@@ -227,15 +243,16 @@
         /// 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 AllocateNewAligned(nuint payloadSize, nuint alignment)
+        private IntPtr AllocateNew(nuint payloadSize)
         {
-            void* raw = _allocator.AllocAligned(payloadSize, alignment);
-            _totalReserved += payloadSize;
+            nuint total = payloadSize + (nuint)sizeof(BlockHeader);
+            void* raw = _allocator.Alloc(total);
+            _totalReserved += total;
             _totalAllocations++;
+
             return (IntPtr)raw;
         }
 
@@ -247,154 +264,100 @@
         /// <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;
-
+                void* ptr = null;
                 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>
-        /// 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;
+                    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);
             }
-            return mIdx;
         }
 
-        /// <summary>
-        /// Allocates a small-size block using bucketed free lists.
-        /// </summary>
-        private void* AllocateSmallAligned(nuint size, nuint alignment, bool zero)
+        /// <summary>Allocates a small-size block using bucketed free lists.</summary>
+        private void* AllocateSmall(nuint size, bool zero)
         {
             nuint bucket = GetSizeClass(size);
-            List<AllocationMetadata> free = _smallFree[bucket];
+            Stack<IntPtr> stack = _smallFree[bucket];
 
-            // Attempt to find an index match the requested size
-            int mIdx = FindIndexAlignmentMatch(alignment, free);
+            IntPtr block = stack.Count > 0
+                ? stack.Pop()
+                : AllocateNew(bucket);
 
-            if (mIdx >= 0)
-            {
-                AllocationMetadata meta = free[mIdx];
-                free.RemoveAt(mIdx);
+            BlockHeader* header = (BlockHeader*)block;
+            header->Size = bucket;
 
-                _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(nUser, 0, (uint)bucket);
-            }
-            return nUser;
+                Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
+
+            return user;
         }
 
         /// <summary>Allocates a medium-size block using best-fit reuse.</summary>
-        private void* AllocateMediumAligned(nuint size, nuint alignment, bool zero)
+        private void* AllocateMedium(nuint size, bool zero)
         {
-            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _mediumFree)
+            foreach (var kv in _mediumFree)
             {
                 if (kv.Key >= size && kv.Value.Count > 0)
                 {
-                    List<AllocationMetadata> free = kv.Value;
+                    var block = kv.Value.Pop();
+                    var header = (BlockHeader*)block;
+                    _totalInUse += header->Size;
 
-                    // Attempt to find an index match the requested size
-                    int mIdx = FindIndexAlignmentMatch(alignment, free);
+                    void* user = header + 1;
+                    if (zero)
+                        Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
 
-                    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;
-                    }
+                    return user;
                 }
             }
-            // no match
-            IntPtr newBlock = AllocateNewAligned(size, alignment);
+
+            var newBlock = AllocateNew(size);
+            var newHeader = (BlockHeader*)newBlock;
+            newHeader->Size = size;
             _totalInUse += size;
 
-            void* nUser = (void*)newBlock;
+            void* newUser = newHeader + 1;
             if (zero)
             {
-                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
-            }               
+                Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
+            }
+                
 
-            return nUser;
+            return newUser;
         }
 
-        /// <summary>
-        /// Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).
-        /// </summary>
-        private void* AllocateLargeAligned(nuint size, nuint alignment, bool zero)
+        /// <summary>Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).</summary>
+        private void* AllocateLarge(nuint size, bool zero)
         {
-            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _largeFree)
+            foreach (var kv in _largeFree)
             {
                 nuint blockSize = kv.Key;
                 if (blockSize < size || kv.Value.Count == 0)
@@ -402,46 +365,41 @@
                     continue;
                 }
 
-                // check waste tolerance before attempting to find an alignment match
                 nuint waste = blockSize - size;
-                bool acceptable = waste <= _largeMaxWasteBytes || ((double)blockSize / size) <= _largeWasteRatioLimit;
+                bool acceptable =
+                    waste <= _largeMaxWasteBytes ||
+                    ((double)blockSize / size) <= _largeWasteRatioLimit;
 
                 if (!acceptable)
                 {
                     continue;
                 }
 
-                List<AllocationMetadata> free = kv.Value;
+                var block = kv.Value.Pop();
+                var header = (BlockHeader*)block;
+                _totalInUse += header->Size;
 
-                // Attempt to find an index match the requested size
-                int mIdx = FindIndexAlignmentMatch(alignment, free);
-
-                if (mIdx >= 0)
+                void* user = header + 1;
+                if (zero)
                 {
-                    AllocationMetadata meta = free[mIdx];
-                    free.RemoveAt(mIdx);
-
-                    _totalInUse += blockSize;
-                    void* user = (void*)meta.Block;
-                    if (zero)
-                    {
-                        Unsafe.InitBlockUnaligned(user, 0, (uint)blockSize);
-                    }
-
-                    return user;
+                    Unsafe.InitBlockUnaligned(user, 0, (uint)header->Size);
                 }
+
+                return user;
             }
 
-            // no match
-
-            IntPtr block = AllocateNewAligned(size, alignment);
+            var newBlock = AllocateNew(size);
+            var newHeader = (BlockHeader*)newBlock;
+            newHeader->Size = size;
             _totalInUse += size;
-            void* nUser = (void*)block;
+
+            void* newUser = newHeader + 1;
             if (zero)
             {
-                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
+                Unsafe.InitBlockUnaligned(newUser, 0, (uint)size);
             }
-            return nUser;
+
+            return newUser;
         }
 
         /// <summary>
@@ -463,41 +421,23 @@
 
             lock (_lock)
             {
-                nuint size = handle.ByteCount;
-                nuint alignment = handle.Alignment;
+                var header = ((BlockHeader*)handle.Pointer) - 1;
+                nuint size = header->Size;
                 _totalInUse -= size;
 
                 if (size <= _smallThreshold)
-                {
-                    _smallFree[size].Add(new AllocationMetadata
-                    {
-                        Block = (IntPtr)handle.Pointer,
-                        Alignment = alignment
-                    });
-                }
+                    _smallFree[size].Push((IntPtr)header);
                 else if (size <= _mediumThreshold)
                 {
-                    if (!_mediumFree.TryGetValue(size, out List<AllocationMetadata>? free))
-                    {
-                        _mediumFree[size] = free = [];
-                    }
-                    free.Add(new AllocationMetadata
-                    {
-                        Block = (IntPtr)handle.Pointer,
-                        Alignment = alignment
-                    });
+                    if (!_mediumFree.TryGetValue(size, out var stack))
+                        _mediumFree[size] = stack = new Stack<IntPtr>();
+                    stack.Push((IntPtr)header);
                 }
                 else
                 {
-                    if (!_largeFree.TryGetValue(size, out List<AllocationMetadata>? free))
-                    {
-                        _largeFree[size] = free = [];
-                    }
-                    free.Add(new AllocationMetadata
-                    {
-                        Block = (IntPtr)handle.Pointer,
-                        Alignment = alignment
-                    });
+                    if (!_largeFree.TryGetValue(size, out var stack))
+                        _largeFree[size] = stack = new Stack<IntPtr>();
+                    stack.Push((IntPtr)header);
                 }
             }
         }
@@ -517,17 +457,17 @@
         }
 
         /// <summary>Helper to free all blocks in a dictionary of free stacks.</summary>
-        private void PruneDictionary(IDictionary<nuint, List<AllocationMetadata>> dict)
+        private void PruneDictionary(IDictionary<nuint, Stack<IntPtr>> dict)
         {
-            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in dict)
+            foreach (var kv in dict)
             {
-                List<AllocationMetadata> list = kv.Value;
-                while (list.Count > 0)
+                var stack = kv.Value;
+                while (stack.Count > 0)
                 {
-                    var item = list[^1];
-                    list.RemoveAt(list.Count - 1);
-                    _allocator.FreeAligned((void*)item.Block, item.Alignment);
-                    _totalReserved -= kv.Key;
+                    var block = stack.Pop();
+                    var header = (BlockHeader*)block;
+                    _allocator.Free((void*)block);
+                    _totalReserved -= (header->Size + (nuint)sizeof(BlockHeader));
                 }
             }
         }