Feature/memory arena improvements (#633)

* First pass at MemoryArena.

* First pass at RandomGenerator.

* Extract TestContext into external source file.

* Fix warning on printf.

* Use enum classes for Test enums.
OutputMode -> TestOutputMode.

* First pass at FreeList unit test.

* Auto registering tests.
Improvements to RandomGenerator.

* Remove the need for unitTest headers - cos can use registering.

* Added unitTest for MemoryArena.

* Do unit tests.

* Fix typo.

* Fix problem limiting errors from TestContext.

* Refactor of MemoryArena
* Removed the ability to rewind (to improve memory usage/simplify)
* Better memory usage - around oversized blocks
  + Will keep allocating from a normal block if more than 1/3 memory left, or an oversided block is allocated

* Better unitTest coverage for MemoryArena.

* Fixes based on code review

* Remove e prefix from enum class types for TestContext
* Added extra checking for allocations sizes
* Fixed some typos
* Added std::is_pod test to allocateAndCopyArray

* Add include for is_pod needed for linux build.

2 files changed, 274 insertions, 233 deletions

diff --git a/source/core/slang-memory-arena.cpp b/source/core/slang-memory-arena.cpp
index bdfd87602..4d8b2ed2a 100644
--- a/source/core/slang-memory-arena.cpp
+++ b/source/core/slang-memory-arena.cpp
@@ -7,32 +7,34 @@ MemoryArena::MemoryArena()
 {
     // Mark as invalid so any alloc call will fail
     m_blockAlignment = 0;
-    m_blockSize = 0;
+    m_blockAllocSize = 0;
 
     // Set up as empty
-    m_blocks = nullptr;
-    _setCurrentBlock(nullptr);
+    m_usedOddBlocks = nullptr;
+    m_availableBlocks = nullptr;
+
+    _resetCurrentBlock();
+
     m_blockFreeList.init(sizeof(Block), sizeof(void*), 16);
 }
 
 MemoryArena::~MemoryArena()
 {
-    _deallocateBlocks();
+    reset();
 }
 
-MemoryArena::MemoryArena(size_t blockSize, size_t blockAlignment)
+MemoryArena::MemoryArena(size_t blockPayloadSize, size_t blockAlignment)
 {
-    _initialize(blockSize, blockAlignment);
+    _initialize(blockPayloadSize, blockAlignment);
 }
 
-void MemoryArena::init(size_t blockSize, size_t blockAlignment)
+void MemoryArena::init(size_t blockPayloadSize, size_t blockAlignment)
 {
-    _deallocateBlocks();
-    m_blockFreeList.reset();
-    _initialize(blockSize, blockAlignment);
+    reset();
+    _initialize(blockPayloadSize, blockAlignment);
 }
   
-void MemoryArena::_initialize(size_t blockSize, size_t alignment)
+void MemoryArena::_initialize(size_t blockPayloadSize, size_t alignment)
 {
     // Alignment must be a power of 2
     assert(((alignment - 1) & alignment) == 0);
@@ -40,99 +42,156 @@ void MemoryArena::_initialize(size_t blockSize, size_t alignment)
     // Must be at least sizeof(void*) in size, as that is the minimum the backing allocator will be
     alignment = (alignment < kMinAlignment) ? kMinAlignment : alignment;
 
+    m_blockPayloadSize = blockPayloadSize;
+    size_t blockAllocSize = blockPayloadSize;
+
     // If alignment required is larger then the backing allocators then
     // make larger to ensure when alignment correction takes place it will be aligned
     if (alignment > kMinAlignment)
     {
-        blockSize += alignment;
+        blockAllocSize += alignment;
     }
 
-    m_blockSize = blockSize;
+    m_blockAllocSize = blockAllocSize;
     m_blockAlignment = alignment;
-    m_blocks = nullptr;
-    _setCurrentBlock(nullptr);
+
+    m_availableBlocks = nullptr;
+    m_usedOddBlocks = nullptr;
+
     m_blockFreeList.init(sizeof(Block), sizeof(void*), 16);
+
+    _resetCurrentBlock();
 }
 
-void* MemoryArena::_allocateAligned(size_t size, size_t alignment)
+void MemoryArena::_resetCurrentBlock()
 {
-    assert(size);
-    // Can't be space in the current block -> so we can either place in next, or in a new block
-    _newCurrentBlock(size, alignment);
-    uint8_t* const current = m_current;
-    // If everything has gone to plan, must be space here...
-    assert(current + size <= m_end);
-    m_current = current + size;
-    return current;
+    m_start = nullptr;
+    m_end = nullptr;
+    m_current = nullptr;
+
+    m_usedBlocks = nullptr;
+}
+
+void MemoryArena::_addCurrentBlock(Block* block)
+{
+    // Set up for allocation from
+    m_end = block->m_end;
+    m_start = block->m_start;
+    m_current = m_start;
+
+    // Add to linked list of used block, making it the top used block
+    block->m_next = m_usedBlocks;
+    m_usedBlocks = block;
 }
 
-void MemoryArena::_deallocateBlocks()
+void MemoryArena::_deallocateBlocksPayload(Block* start)
 {
-    Block* currentBlock = m_blocks;
-    while (currentBlock)
+    Block* cur = start;
+    while (cur)
     {
         // Deallocate the block
-        ::free(currentBlock->m_alloc);
-        // next block
-        currentBlock = currentBlock->m_next;
+        ::free(cur->m_alloc);
+        cur = cur->m_next;
     }
-    // Can deallocate all blocks to
-    m_blockFreeList.deallocateAll();
 }
 
-void MemoryArena::_setCurrentBlock(Block* block)
+void MemoryArena::_deallocateBlocks(Block* start)
 {
-    if (block)
+    Block* cur = start;
+    while (cur)
     {
-        m_end = block->m_end;
-        m_start = block->m_start;
-        m_current = m_start;
+        Block* next = cur->m_next;
+        // Deallocate the block
+        ::free(cur->m_alloc);
+
+        m_blockFreeList.deallocate(cur);
+        cur = next;
     }
-    else
+}
+
+/* static */MemoryArena::Block* MemoryArena::_joinBlocks(Block* pre, Block* post)
+{
+    if (pre && post)
     {
-        m_start = nullptr;
-        m_end = nullptr;
-        m_current = nullptr;
+        // If both are actual lists, concat post at end of pre
+        Block* cur = pre;
+        while (cur->m_next)
+        {
+            cur = cur->m_next;
+        }
+        // Attach post to end of pre
+        cur->m_next = post;
     }
-    m_currentBlock = block;
+    return pre ? pre : post;
 }
 
-MemoryArena::Block* MemoryArena::_newCurrentBlock(size_t size, size_t alignment)
+void MemoryArena::deallocateAll()
 {
-    // Make sure init has been called (or has been set up in parameterized constructor)
-    assert(m_blockSize > 0);
-    // Alignment must be a power of 2
-    assert(((alignment - 1) & alignment) == 0);
+    // Free all oversized
+    _deallocateBlocks(m_usedOddBlocks);
+    m_usedOddBlocks = nullptr;
+    // We want to put used blocks onto the start of the available list
+    m_availableBlocks = _joinBlocks(m_usedBlocks, m_availableBlocks);
 
-    // Alignment must at a minimum be block alignment (such if reused the constraints hold)
-    alignment = (alignment < m_blockAlignment) ? m_blockAlignment : alignment;
+    // Reset current block
+    _resetCurrentBlock();
+}
 
-    const size_t alignMask = alignment - 1;
+void MemoryArena::reset()
+{
+    _deallocateBlocksPayload(m_usedOddBlocks);
+    _deallocateBlocksPayload(m_usedBlocks);
+    _deallocateBlocksPayload(m_availableBlocks);
+
+    m_blockFreeList.reset();
+
+    m_usedOddBlocks = nullptr;
+    m_availableBlocks = nullptr;
 
-    // First try the next block (if there is one)
+    _resetCurrentBlock();
+}
+
+const MemoryArena::Block* MemoryArena::_findNonCurrent(const void* data, size_t size) const
+{
+    // It must either be m_usedOversizedBlocks or after m_usedBlocks (because m_usedBlocks is m_current)
+    const Block* block = _findInBlocks(m_usedOddBlocks, data, size);
+    if (block) 
     {
-        Block* next = m_currentBlock ? m_currentBlock->m_next : m_blocks;
-        if (next)
+        return block;
+    }
+    return m_usedBlocks ? _findInBlocks(m_usedBlocks->m_next, data, size) : nullptr;
+}
+
+const MemoryArena::Block* MemoryArena::_findInBlocks(const Block* block, const void* data, size_t size) const
+{
+    const uint8_t* ptr = (const uint8_t*)data;
+    while (block)
+    {
+        if (ptr >= block->m_start && ptr + size <= block->m_end)
         {
-            // Align could be done from the actual allocation start, but doing so would mean a pointer which
-            // didn't hit the constraint of being between start/end
-            // So have to align conservatively using start
-            uint8_t* memory = (uint8_t*)((size_t(next->m_start) + alignMask) & ~alignMask);
-
-            // Check if can fit block in
-            if (memory + size <= next->m_end)
-            {
-                _setCurrentBlock(next);
-                return next;
-            }
+            return block;
         }
+        block = block->m_next;
     }
+    return nullptr;
+}
 
-    // The size of the block must be at least large enough to take into account alignment
-    size_t allocSize = (alignment <= kMinAlignment) ? size : (size + alignment);
+MemoryArena::Block* MemoryArena::_newNormalBlock()
+{
+    if (m_availableBlocks)
+    {
+        // We have an available block..
+        Block* block = m_availableBlocks;
+        m_availableBlocks = block->m_next;
+        return block;
+    }
 
-    // The minimum block size should be at least m_blockSize
-    allocSize = (allocSize < m_blockSize) ? m_blockSize : allocSize;
+    return _newBlock(m_blockAllocSize, m_blockAlignment);
+}
+
+MemoryArena::Block* MemoryArena::_newBlock(size_t allocSize, size_t alignment)
+{
+    assert(alignment >= m_blockAlignment);
 
     // Allocate block
     Block* block = (Block*)m_blockFreeList.allocate();
@@ -140,6 +199,7 @@ MemoryArena::Block* MemoryArena::_newCurrentBlock(size_t size, size_t alignment)
     {
         return nullptr;
     }
+
     // Allocate the memory
     uint8_t* alloc = (uint8_t*)::malloc(allocSize);
     if (!alloc)
@@ -147,6 +207,9 @@ MemoryArena::Block* MemoryArena::_newCurrentBlock(size_t size, size_t alignment)
         m_blockFreeList.deallocate(block);
         return nullptr;
     }
+
+    const size_t alignMask = alignment - 1;
+
     // Do the alignment on the allocation
     uint8_t* const start = (uint8_t*)((size_t(alloc) + alignMask) & ~alignMask);
 
@@ -156,89 +219,103 @@ MemoryArena::Block* MemoryArena::_newCurrentBlock(size_t size, size_t alignment)
     block->m_end = alloc + allocSize;
     block->m_next = nullptr;
 
-    // Insert block into list
-    if (m_currentBlock)
-    {
-        // Insert after current block
-        block->m_next = m_currentBlock->m_next;
-        m_currentBlock->m_next = block;
-    }
-    else
-    {
-        // Add to start of the list of the blocks
-        block->m_next = m_blocks;
-        m_blocks = block;
-    }
-    _setCurrentBlock(block);
     return block;
 }
 
-MemoryArena::Block* MemoryArena::_findBlock(const void* alloc, Block* endBlock) const
+void* MemoryArena::_allocateAlignedFromNewBlock(size_t size, size_t alignment)
 {
-    const uint8_t* ptr = (const uint8_t*)alloc;
+    // Make sure init has been called (or has been set up in parameterized constructor)
+    assert(m_blockAllocSize > 0);
+    // Alignment must be a power of 2
+    assert(((alignment - 1) & alignment) == 0);
+
+    // Alignment must at a minimum be block alignment (such if reused the constraints hold)
+    alignment = (alignment < m_blockAlignment) ? m_blockAlignment : alignment;
+
+    const size_t alignMask = alignment - 1;
 
-    Block* block = m_blocks;
-    while (block != endBlock)
+    // The size of the block must be at least large enough to take into account alignment
+    size_t allocSize = (alignment <= kMinAlignment) ? size : (size + alignment);
+
+    const size_t currentRemainSize = size_t(m_end - m_current);
+
+    // a) Allocate a new normal block and make current
+    // b) Allocate a new normal block (leave current)
+    // c) Allocate a new 'oversized' block (leave current)
+    // 
+    // We only bother with a and c here, as b is only really usable if size is close to m_blockAllocSize
+
+    // If there is > 1/3 of block remaining, or the block required is too big to fit use an 'oversized' block
+    if ((currentRemainSize * 3 > m_blockPayloadSize) || (allocSize > m_blockAllocSize))
     {
-        if (ptr >= block->m_start && ptr < block->m_end)
+        // We don't change current because there is enough remaining
+        Block* block = _newBlock(allocSize, alignment);
+        if (!block)
         {
-            return block;
+            return nullptr;
         }
-        block = block->m_next;
+        // Add to odd used blocks list
+        block->m_next = m_usedOddBlocks;
+        m_usedOddBlocks = block;
+        // NOTE! We are keeping the previous m_current current, so any remaining space can still be used
+        // Return the address
+        return block->m_start;
     }
-    return nullptr;
+
+    // Must be allocatable within a normal block
+    assert(allocSize <= m_blockAllocSize);
+    Block* block = _newNormalBlock();
+    if (!block)
+    {
+        return nullptr;
+    }
+    // It's a new regular block...
+    _addCurrentBlock(block);
+
+    // Do the aligned allocation (which must fit) by aligning the pointer
+    uint8_t* memory = (uint8_t*)((size_t(m_current) + alignMask) & ~alignMask);
+    // It must fit if the previous code is correct...
+    assert(memory + size <= m_end);
+    // Move the current pointer
+    m_current = memory + size;
+    return memory;
 }
 
-MemoryArena::Block* MemoryArena::_findPreviousBlock(Block* block)
+size_t MemoryArena::_calcBlocksUsedMemory(const Block* block) const
 {
-    Block* currentBlock = m_blocks;
-    while (currentBlock)
+    size_t total = 0;
+    while (block)
     {
-        if (currentBlock->m_next == block)
-        {
-            return currentBlock;
-        }
-        currentBlock = currentBlock->m_next;
+        total += size_t(block->m_end - block->m_start);
+        block = block->m_next;
     }
-    return nullptr;
+    return total;
 }
 
-void MemoryArena::deallocateAll()
+size_t MemoryArena::_calcBlocksAllocatedMemory(const Block* block) const
 {
-    Block** prev = &m_blocks;
-    Block* block = m_blocks;
-
+    size_t total = 0;
     while (block)
     {
-        if (size_t(block->m_end - block->m_alloc) > m_blockSize)
-        {
-            // Oversized block so we need to free it and remove from the list
-            Block* nextBlock = block->m_next;
-            *prev = nextBlock;
-            // Free the backing memory
-            ::free(block->m_alloc);
-            // Free the block
-            m_blockFreeList.deallocate(block);
-            // prev stays the same, now working on next tho
-            block = nextBlock;
-        }
-        else
-        {
-            // Onto next
-            prev = &block->m_next;
-            block = block->m_next;
-        }
+        total += size_t(block->m_end - block->m_alloc);
+        block = block->m_next;
     }
+    return total;
+}
 
-    // Make the first current (if any)
-    _setCurrentBlock(m_blocks);
+size_t MemoryArena::calcTotalMemoryUsed() const
+{
+    return _calcBlocksUsedMemory(m_usedOddBlocks)
+        + (m_usedBlocks ? _calcBlocksUsedMemory(m_usedBlocks->m_next) : 0) +
+        size_t(m_current - m_start);
 }
 
-void MemoryArena::reset()
+size_t MemoryArena::calcTotalMemoryAllocated() const
 {
-    _deallocateBlocks();
-    m_blocks = nullptr;
-    _setCurrentBlock(nullptr);
+    return _calcBlocksAllocatedMemory(m_usedOddBlocks) + 
+        _calcBlocksAllocatedMemory(m_usedBlocks) +
+        _calcBlocksAllocatedMemory(m_availableBlocks);
 }
 
+
 } // namespace Slang
diff --git a/source/core/slang-memory-arena.h b/source/core/slang-memory-arena.h
index acbac51dd..1c59ee7d7 100644
--- a/source/core/slang-memory-arena.h
+++ b/source/core/slang-memory-arena.h
@@ -8,18 +8,16 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <type_traits>
+
 #include "slang-free-list.h"
 
 namespace Slang { 
 
-/** Defines arena allocator where allocations are made very quickly, but that deallocations
-can only be performed in reverse order, or with the client code knowing a previous deallocation (say with
- deallocateAllFrom), automatically deallocates everything after it.
+/** MemoryArena provides provides very fast allocation of small blocks, by aggregating many small allocations 
+over smaller amount of larger blocks. A typical small unaligned allocation is a pointer bump.  
 
-It works by allocating large blocks and then cutting out smaller pieces as requested. If a piece of memory is
-deallocated, it either MUST be in reverse allocation order OR the subsequent allocations are implicitly 
-deallocated too, and therefore accessing their memory is now undefined behavior. Allocations are made 
-contiguously from the current block. If there is no space in the current block, the 
+Allocations are made contiguously from the current block. If there is no space in the current block, the 
 next block (which is unused) if available is checked. If that works, an allocation is made from the next block. 
 If not a new block is allocated that can hold at least the allocation with required alignment.
  
@@ -27,7 +25,7 @@ All memory allocated can be deallocated very quickly and without a client having
 All memory allocated will be freed on destruction - or with reset.
 
 A memory arena can have requests larger than the block size. When that happens they will just be allocated
-from the heap. As such 'oversized blocks' are seen as unusual and potentially wasteful they are deallocated
+from the heap. As such 'odd blocks' are seen as unusual and potentially wasteful so they are deallocated
 when deallocateAll is called, whereas regular size blocks will remain allocated for fast subsequent allocation.
 
 It is intentional that blocks information is stored separately from the allocations that store the
@@ -35,18 +33,22 @@ user data. This is so that alignment permitting, block allocations sizes can be
 For large power of 2 backing allocations this might mean a page/pages directly allocated by the OS for example.
 Also means better cache coherency when traversing blocks -> as generally they will be contiguous in memory.
 
-Also note that allocateUnaligned can be used for slightly faster aligned allocations. All blocks allocated internally
+Note that allocateUnaligned can be used for slightly faster aligned allocations. All blocks allocated internally
 are aligned to the blockAlignment passed to the constructor. If subsequent allocations (of any type) sizes are of that
 alignment or larger then no alignment fixing is required (because allocations are contiguous) and so 'allocateUnaligned'
 will return allocations of blockAlignment alignment.
+
+If many 'odd' allocations occur it probably means that the block size should be increased. 
 */
 class MemoryArena
 {
 public:
     typedef MemoryArena ThisType;
 
-    static const size_t kMinAlignment = sizeof(void*);  ///< The minimum alignment of the backing memory allocator.
-
+        /** The minimum alignment of the backing memory allocator.
+        NOTE! That this should not be greater than the alignment of the underlying allocator.
+        */
+    static const size_t kMinAlignment = sizeof(void*);  
         /** Determines if an allocation is consistent with an allocation from this arena.
 
         The test cannot say definitively if this was such an allocation, because the exact details
@@ -78,8 +80,7 @@ public:
 
         /** Allocate some aligned memory of at least size bytes 
         @param size Size of allocation wanted (must be > 0).
-        @param alignment Alignment of allocation - must be a power of 2.
-        @return The allocation (or nullptr if unable to allocate). Will be at least 'alignment' alignment or better. */
+        @return The allocation (or nullptr if unable to allocate).  */
     void* allocateUnaligned(size_t sizeInBytes);
 
         /** Allocates a null terminated string.
@@ -109,12 +110,6 @@ public:
     template <typename T>
     T* allocateAndZeroArray(size_t size);
 
-        /// Deallocate the last allocation. If data is not from the last allocation then the behavior is undefined.
-    void deallocateLast(void* data);
-
-        /// Deallocate this allocation and all remaining after it.
-    void deallocateAllFrom(void* dataStart);
-
         /** Deallocates all allocated memory. That backing memory will generally not be released so
          subsequent allocation will be fast, and from the same memory. Note though that 'oversize' blocks
          will be deallocated. */
@@ -128,10 +123,15 @@ public:
         /// Gets the block alignment that is passed at initialization otherwise 0 an invalid block alignment.
     size_t getBlockAlignment() const { return m_blockAlignment; }
 
+        /// Estimate of total amount of memory used in bytes. The number can never be smaller than actual used memory but may be larger
+    size_t calcTotalMemoryUsed() const;
+        /// Total memory allocated in bytes
+    size_t calcTotalMemoryAllocated() const;
+
         /// Default Ctor
     MemoryArena();
         /// Construct with block size and alignment. Block alignment must be a power of 2.
-    MemoryArena(size_t blockSize, size_t blockAlignment = kMinAlignment);
+    MemoryArena(size_t blockPayloadSize, size_t blockAlignment = kMinAlignment);
 
         /// Dtor
     ~MemoryArena();
@@ -139,32 +139,51 @@ public:
 protected:
     struct Block
     {
-        Block* m_next;
-        uint8_t* m_alloc;
-        uint8_t* m_start;
-        uint8_t* m_end;
+        Block* m_next;                  ///< Singly linked list of blocks
+        uint8_t* m_alloc;               ///< Allocation start (ie what to free)
+        uint8_t* m_start;               ///< Start of payload (takes into account alignment)
+        uint8_t* m_end;                 ///< End of payload (m_start to m_end defines payload)
     };
 
-    void _initialize(size_t blockSize, size_t blockAlignment);
+    void _initialize(size_t blockPayloadSize, size_t blockAlignment);
+
+        /// Delete the linked list of blocks specified by start
+    void _deallocateBlocks(Block* start);
+        /// Delete the linked list of blocks payloads specified by start
+    void _deallocateBlocksPayload(Block* start);
+
+    void _resetCurrentBlock();
+    void _addCurrentBlock(Block* block);
+
+    static Block* _joinBlocks(Block* pre, Block* post);
+
+        /// Create a new block with regular block alignment 
+    Block* _newNormalBlock();
+        /// Allocates a new block with allocSize and alignment 
+    Block* _newBlock(size_t allocSize, size_t alignment);
 
-    void* _allocateAligned(size_t size, size_t alignment);
-    void _deallocateBlocks();
+    void* _allocateAlignedFromNewBlock(size_t size, size_t alignment);
 
-    void _setCurrentBlock(Block* block);
+        /// Find block that contains data/size that is _NOT_ current (ie not first block in m_usedBlocks)
+    const Block* _findNonCurrent(const void* data, size_t size) const;
+    const Block* _findInBlocks(const Block* block, const void* data, size_t size) const;
 
-    Block* _newCurrentBlock(size_t size, size_t alignment);
-    Block* _findBlock(const void* alloc, Block* endBlock = nullptr) const;
-    Block* _findPreviousBlock(Block* block);
+    size_t _calcBlocksUsedMemory(const Block* block) const;
+    size_t _calcBlocksAllocatedMemory(const Block* block) const;
 
-    uint8_t* m_start;
-    uint8_t* m_end;
-    uint8_t* m_current;
-    size_t m_blockSize; 
-    size_t m_blockAlignment;
-    Block* m_blocks;
-    Block* m_currentBlock;
+    uint8_t* m_start;               ///< The start of the current block (pointed to by m_usedBlocks)
+    uint8_t* m_end;                 ///< The end of the current block
+    uint8_t* m_current;             ///< The current position in current block
 
-    FreeList m_blockFreeList;
+    size_t m_blockPayloadSize;      ///< The size of the payload of a block
+    size_t m_blockAllocSize;        ///< The size of a block allocation (must be the same size or bigger than m_blockPayloadSize)
+    size_t m_blockAlignment;        ///< The alignment applied to used blocks
+
+    Block* m_availableBlocks;       ///< Standard sized blocks that are available
+    Block* m_usedBlocks;            ///< List of all normal sized used blocks. The first one is the 'current block'
+    Block* m_usedOddBlocks;         ///< Used 'odd' blocks - blocks can actually be smaller than normal blocks, but are typically larger. 
+
+    FreeList m_blockFreeList;       ///< Holds all of the blocks for fast allocation/free
 
     private:
     // Disable
@@ -176,21 +195,14 @@ protected:
 inline bool MemoryArena::isValid(const void* data, size_t size) const
 {
     assert(size);
-
     uint8_t* ptr = (uint8_t*)data;
-    // Is it in current
-    if (ptr >= m_start && ptr + size <= m_current)
-    {
-        return true;
-    }
-    // Is it in a previous block?
-    Block* block = _findBlock(data, m_currentBlock);
-    return block && (ptr >= block->m_start && (ptr + size) <= block->m_end);
+    return (ptr >= m_start && ptr + size <= m_current) || _findNonCurrent(data, size) != nullptr;
 }
 
 // --------------------------------------------------------------------------
 SLANG_FORCE_INLINE void* MemoryArena::allocateUnaligned(size_t size)
 {
+    assert(size > 0);
     // Align with the minimum alignment
     uint8_t* mem = m_current;
     uint8_t* end = mem + size;
@@ -201,13 +213,14 @@ SLANG_FORCE_INLINE void* MemoryArena::allocateUnaligned(size_t size)
     }
     else
     {
-        return _allocateAligned(size, m_blockAlignment);
+        return _allocateAlignedFromNewBlock(size, kMinAlignment);
     }
 }
 
 // --------------------------------------------------------------------------
 SLANG_FORCE_INLINE void* MemoryArena::allocate(size_t size)
 {
+    assert(size > 0);
     // Align with the minimum alignment
     const size_t alignMask = kMinAlignment - 1;
     uint8_t* mem = (uint8_t*)((size_t(m_current) + alignMask) & ~alignMask);
@@ -219,13 +232,14 @@ SLANG_FORCE_INLINE void* MemoryArena::allocate(size_t size)
     }
     else
     {
-        return _allocateAligned(size, kMinAlignment);
+        return _allocateAlignedFromNewBlock(size, kMinAlignment);
     }
 }
 
 // --------------------------------------------------------------------------
 inline void* MemoryArena::allocateAligned(size_t size, size_t alignment)
 {
+    assert(size > 0);
     // Alignment must be a power of 2
     assert(((alignment - 1) & alignment) == 0);
 
@@ -240,7 +254,7 @@ inline void* MemoryArena::allocateAligned(size_t size, size_t alignment)
     }
     else
     {
-        return _allocateAligned(size, alignment);
+        return _allocateAlignedFromNewBlock(size, alignment);
     }
 }
 
@@ -290,6 +304,8 @@ inline T* MemoryArena::allocateArray(size_t count)
 template <typename T>
 inline T* MemoryArena::allocateAndCopyArray(const T* arr, size_t size)
 {
+    SLANG_COMPILE_TIME_ASSERT(std::is_pod<T>::value);
+
     if (size > 0)
     {
         const size_t totalSize = sizeof(T) * size;
@@ -315,59 +331,6 @@ inline T* MemoryArena::allocateAndZeroArray(size_t size)
 }
 
 // --------------------------------------------------------------------------
-inline void MemoryArena::deallocateLast(void* data)
-{
-    // See if it's in current block
-    uint8_t* ptr = (uint8_t*)data;
-    if (ptr >= m_start && ptr < m_current)
-    {
-        // Then just go back
-        m_current = ptr;
-    }
-    else
-    {
-        // Only called if not in the current block. Therefore can only be in previous
-        Block* prevBlock = _findPreviousBlock(m_currentBlock);
-        if (prevBlock == nullptr || (!(ptr >= prevBlock->m_start && ptr < prevBlock->m_end)))
-        {
-            assert(!"Allocation not found");
-            return;
-        }
-
-        // Make the previous block the current
-        _setCurrentBlock(prevBlock);
-        // Make the current the alloc freed
-        m_current = ptr;
-    }
-}
-
-// --------------------------------------------------------------------------
-inline void MemoryArena::deallocateAllFrom(void* data)
-{
-    // See if it's in current block, and is allocated (ie < m_current)
-    uint8_t* ptr = (uint8_t*)data;
-    if (ptr >= m_start && ptr < m_current)
-    {
-        // If it's in current block, then just go back
-        m_current = ptr;
-        return;
-    }
-
-    // Search all blocks prior to current block
-    Block* block = _findBlock(data, m_currentBlock);
-    assert(block);
-    if (!block)
-    {
-        return;
-    }
-    // Make this current block
-    _setCurrentBlock(block);
-
-    // Move the pointer to the allocations position
-    m_current = ptr;
-}
-
-// --------------------------------------------------------------------------
 inline void MemoryArena::adjustToBlockAlignment()
 {
     const size_t alignMask = m_blockAlignment - 1;
@@ -377,8 +340,9 @@ inline void MemoryArena::adjustToBlockAlignment()
     // This test could be avoided if we aligned m_end, but depending on block alignment that might waste some space
     if (ptr > m_end)
     {
-        // We'll need a new block to make this alignment
-        _newCurrentBlock(0, m_blockAlignment);
+        // We'll need a new block to make this alignment. Allocate a byte, and then rewind it.
+        _allocateAlignedFromNewBlock(1, 1);
+        m_current = m_usedBlocks->m_start; 
     }
     else
     {