MemoryArena rewindability/Improved IRInst construction (#837)

* Make MemoryArena rewindable.

* Add test for rewinding MemoryArena

* Use memory rewinding in IRInst lookup instead of malloc/free.

* Small tidy.

* Don't bother recreating instruction if after lookup it's found it's new.

* Fix 32 bit signed compare issue.

4 files changed, 240 insertions, 85 deletions

diff --git a/source/core/slang-memory-arena.cpp b/source/core/slang-memory-arena.cpp
index a0a89155c..407c6eedd 100644
--- a/source/core/slang-memory-arena.cpp
+++ b/source/core/slang-memory-arena.cpp
@@ -10,7 +10,7 @@ MemoryArena::MemoryArena()
     m_blockAllocSize = 0;
 
     // Set up as empty
-    m_usedOddBlocks = nullptr;
+    m_usedBlocks = nullptr;
     m_availableBlocks = nullptr;
 
     _resetCurrentBlock();
@@ -56,8 +56,7 @@ void MemoryArena::_initialize(size_t blockPayloadSize, size_t alignment)
     m_blockAlignment = alignment;
 
     m_availableBlocks = nullptr;
-    m_usedOddBlocks = nullptr;
-
+    
     m_blockFreeList.init(sizeof(Block), sizeof(void*), 16);
 
     _resetCurrentBlock();
@@ -84,6 +83,16 @@ void MemoryArena::_addCurrentBlock(Block* block)
     m_usedBlocks = block;
 }
 
+void MemoryArena::_setCurrentBlock(Block* block)
+{
+    // Set up for allocation from
+    m_end = block->m_end;
+    m_start = block->m_start;
+    m_current = m_start;
+
+    assert(m_usedBlocks == block);
+}
+
 void MemoryArena::_deallocateBlocksPayload(Block* start)
 {
     Block* cur = start;
@@ -109,60 +118,68 @@ void MemoryArena::_deallocateBlocks(Block* start)
     }
 }
 
-/* static */MemoryArena::Block* MemoryArena::_joinBlocks(Block* pre, Block* post)
+bool MemoryArena::_isNormalBlock(Block* block)
 {
-    if (pre && post)
+    size_t blockSize = size_t(block->m_end - block->m_start);
+    return (blockSize == m_blockAllocSize) && ((size_t(block->m_start) & (m_blockAlignment - 1)) == 0);
+}
+
+void MemoryArena::_deallocateBlock(Block* block)
+{
+    // If it's a normal block then make it available
+    if (_isNormalBlock(block))
     {
-        // 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;
+        block->m_next = m_availableBlocks;
+        m_availableBlocks = block;
+    }
+    else
+    {
+        // Must be oversized so free it
+        ::free(block->m_alloc);
+        // Free it in the block list
+        m_blockFreeList.deallocate(block);
     }
-    return pre ? pre : post;
 }
 
 void MemoryArena::deallocateAll()
 {
-    // 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);
+    // we need to rewind through m_usedBlocks -> seeing it the are normal sized or not
 
+    Block* block = m_usedBlocks;
+    while (block)
+    {
+        Block* next = block->m_next;
+        _deallocateBlock(block);
+        block = next;
+    }
+    
     // Reset current block
     _resetCurrentBlock();
 }
 
 void MemoryArena::reset()
 {
-    _deallocateBlocksPayload(m_usedOddBlocks);
     _deallocateBlocksPayload(m_usedBlocks);
     _deallocateBlocksPayload(m_availableBlocks);
 
     m_blockFreeList.reset();
 
-    m_usedOddBlocks = nullptr;
     m_availableBlocks = nullptr;
 
     _resetCurrentBlock();
 }
 
-const MemoryArena::Block* MemoryArena::_findNonCurrent(const void* data, size_t size) 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) 
-    {
-        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
+MemoryArena::Block* MemoryArena::_findNonCurrent(const void* data) const
+{
+    return m_usedBlocks ? _findInBlocks(m_usedBlocks->m_next, data) : nullptr;
+}
+
+MemoryArena::Block* MemoryArena::_findInBlocks(Block* block, const void* data, size_t size) const
 {
     const uint8_t* ptr = (const uint8_t*)data;
     while (block)
@@ -176,6 +193,20 @@ const MemoryArena::Block* MemoryArena::_findInBlocks(const Block* block, const v
     return nullptr;
 }
 
+MemoryArena::Block* MemoryArena::_findInBlocks(Block* block, const void* data) const
+{
+    const uint8_t* ptr = (const uint8_t*)data;
+    while (block)
+    {
+        if (ptr >= block->m_start && ptr <= block->m_end)
+        {
+            return block;
+        }
+        block = block->m_next;
+    }
+    return nullptr;
+}
+
 MemoryArena::Block* MemoryArena::_newNormalBlock()
 {
     if (m_availableBlocks)
@@ -249,32 +280,30 @@ void* MemoryArena::_allocateAlignedFromNewBlock(size_t size, size_t alignment)
 
     const size_t currentRemainSize = size_t(m_end - m_current);
 
+    Block* block;
+
+    // There are two scenario
     // 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)
+    // b) Allocate a new 'oversized' block and make current
     // 
-    // We only bother with a and c here, as b is only really usable if size is close to m_blockAllocSize
+    // That by always allocating a new block if oversized, we lose more efficiency in terms of storage (the previous block
+    // may not have been used much). BUT doing so makes it easy to rewind - as the blocks are always in order of allocation.
+    // An improvement might be to have some abstraction that sits on top that can do this tracking (or have the blocks
+    // themselves record if they alias over a previously used block - but we don't bother with this here.
 
     // 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))
     {
-        // We don't change current because there is enough remaining
-        Block* block = _newBlock(allocSize, alignment);
-        if (!block)
-        {
-            return nullptr;
-        }
-        // 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;
+        // This is an oversized block so just allocate the whole thing
+        block = _newBlock(allocSize, alignment);
     }
-
-    // Must be allocatable within a normal block
-    assert(allocSize <= m_blockAllocSize);
-    Block* block = _newNormalBlock();
+    else
+    {
+        // Must be allocatable within a normal block
+        assert(allocSize <= m_blockAllocSize);
+        block = _newNormalBlock();
+    }
+    // If not allocated we are done
     if (!block)
     {
         return nullptr;
@@ -313,17 +342,57 @@ size_t MemoryArena::_calcBlocksAllocatedMemory(const Block* block) const
     return total;
 }
 
+void MemoryArena::_rewindToCursor(const void* cursorIn)
+{
+    // If it's nullptr, then there are no allocation so free all
+    if (cursorIn == nullptr)
+    {
+        deallocateAll();
+        return;
+    }
+   
+    // Find the block that contains the allocation
+    Block* cursorBlock = _findNonCurrent(cursorIn);
+    assert(cursorBlock);
+    if (!cursorBlock)
+    {
+        // If not found it means this address is NOT part any of the active used heap!
+        // Probably an invalid cursor 
+        return;
+    }
+
+    // Deallocate all of the blocks up to the cursor block
+    {
+        Block* block = m_usedBlocks;
+        while (block != cursorBlock)
+        {
+            Block* next = block->m_next;
+            _deallocateBlock(block);
+            block = next;
+        }
+    }
+
+    // The cursor block is now the current block
+    m_usedBlocks = cursorBlock; 
+    _setCurrentBlock(cursorBlock);
+
+    const uint8_t* cursor = (const uint8_t*)cursorIn;
+    // Must be in the range of the currently set block
+    assert(cursor >= m_start && cursor <= m_end);
+
+    // Set the current position where the cursor is
+    m_current = const_cast<uint8_t*>(cursor);
+}
+
 size_t MemoryArena::calcTotalMemoryUsed() const
 {
-    return _calcBlocksUsedMemory(m_usedOddBlocks)
-        + (m_usedBlocks ? _calcBlocksUsedMemory(m_usedBlocks->m_next) : 0) +
+    return (m_usedBlocks ? _calcBlocksUsedMemory(m_usedBlocks->m_next) : 0) +
         size_t(m_current - m_start);
 }
 
 size_t MemoryArena::calcTotalMemoryAllocated() const
 {
-    return _calcBlocksAllocatedMemory(m_usedOddBlocks) + 
-        _calcBlocksAllocatedMemory(m_usedBlocks) +
+    return _calcBlocksAllocatedMemory(m_usedBlocks) +
         _calcBlocksAllocatedMemory(m_availableBlocks);
 }
 
diff --git a/source/core/slang-memory-arena.h b/source/core/slang-memory-arena.h
index 54b380b32..14eca9e02 100644
--- a/source/core/slang-memory-arena.h
+++ b/source/core/slang-memory-arena.h
@@ -133,6 +133,11 @@ public:
         /// Total memory allocated in bytes
     size_t calcTotalMemoryAllocated() const;
 
+        /// Get thue current allocation cursor (memory address where subsequent allocations will be placed if space within the current block)
+    void* getCursor() const { return m_current; }
+        /// Rewind (and effectively deallocate) all allocations *after* the cursor
+    void rewindToCursor(const void* cursor);
+
         /// Default Ctor
     MemoryArena();
         /// Construct with block size and alignment. Block alignment must be a power of 2.
@@ -159,8 +164,9 @@ protected:
 
     void _resetCurrentBlock();
     void _addCurrentBlock(Block* block);
+    void _setCurrentBlock(Block* block);
 
-    static Block* _joinBlocks(Block* pre, Block* post);
+    void _deallocateBlock(Block* block);
 
         /// Create a new block with regular block alignment 
     Block* _newNormalBlock();
@@ -171,11 +177,20 @@ protected:
     void* _allocateAlignedFromNewBlockAndZero(size_t sizeInBytes, size_t alignment);
 
         /// 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 sizeInBytes) const;
-    const Block* _findInBlocks(const Block* block, const void* data, size_t sizeInBytes) const;
+    Block* _findNonCurrent(const void* data, size_t sizeInBytes) const;
+    Block* _findNonCurrent(const void* data) const;
+
+        /// Find a block that contains data starting from block. Returns null ptr if not found
+    Block* _findInBlocks(Block* block, const void* data) const;
+    Block* _findInBlocks(Block* block, const void* data, size_t sizeInBytes) const;
 
     size_t _calcBlocksUsedMemory(const Block* block) const;
     size_t _calcBlocksAllocatedMemory(const Block* block) const;
+        /// Returns true if block can be classed as normal (right size and same or better alignment)
+    bool _isNormalBlock(Block* block);
+
+        /// Handles the rewinding of the cursor for the more complicated cases
+    void _rewindToCursor(const void* cursor);
 
     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
@@ -186,9 +201,8 @@ protected:
     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. 
-
+    Block* m_usedBlocks;            ///< Singly linked list of used blocks. The first one is the 'current block' and m_next is the previously allocated blocks. nullptr terminated.
+    
     FreeList m_blockFreeList;       ///< Holds all of the blocks for fast allocation/free
 
     private:
@@ -377,7 +391,21 @@ inline void MemoryArena::adjustToBlockAlignment()
     }
     assert(size_t(m_current) & alignMask);
 }
+// --------------------------------------------------------------------------
+SLANG_FORCE_INLINE void MemoryArena::rewindToCursor(const void* cursor)
+{
+    // Is it in the current block?
+    {
+        const uint8_t* cur = (const uint8_t*)cursor;
+        if (cur >= m_start && cur <= m_current)
+        {
+            m_current = const_cast<uint8_t*>(cur);
+            return;
+        }
+    }
+    _rewindToCursor(cursor);
+}
 
 } // namespace Slang
 
-#endif // SLANG_MEMORY_ARENA_H
\ No newline at end of file
+#endif // SLANG_MEMORY_ARENA_H
diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index 0a5b8491c..bed506520 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -1458,54 +1458,91 @@ namespace Slang
             operandCount += listOperandCounts[ii];
         }
 
+        // A more aggressive Impl could create the instruction once, and only free the memory
+        // allocated to construct if it was found that the instruction already exists.
+        // Here we just avoid using malloc and use the memoryArena as a cheap way to allocate some
+        // temporary memory.
+        auto& memoryArena = builder->getModule()->memoryArena;
+        void* cursor = memoryArena.getCursor();
+
         // We are going to create a dummy instruction on the stack,
         // which will be used as a key for lookup, so see if we
         // already have an equivalent instruction available to use.
 
         size_t keySize = sizeof(IRInst) + operandCount * sizeof(IRUse);
-        IRInst* keyInst = (IRInst*) malloc(keySize);
-        memset(keyInst, 0, keySize);
+        IRInst* inst = (IRInst*) memoryArena.allocateAndZero(keySize);
+        
+        void* endCursor = memoryArena.getCursor();
+        // Mark as 'unused' cos it is unused on release builds. 
+        SLANG_UNUSED(endCursor);
 
-        new(keyInst) IRInst();
-        keyInst->op = op;
-        keyInst->typeUse.usedValue = type;
-        keyInst->operandCount = (uint32_t) operandCount;
+        new(inst) IRInst();
+        inst->op = op;
+        inst->typeUse.usedValue = type;
+        inst->operandCount = (uint32_t) operandCount;
 
-        IRUse* operand = keyInst->getOperands();
-        for (UInt ii = 0; ii < operandListCount; ++ii)
+        // Don't link up as we may free (if we already have this)
         {
-            UInt listOperandCount = listOperandCounts[ii];
-            for (UInt jj = 0; jj < listOperandCount; ++jj)
+            IRUse* operand = inst->getOperands();
+            for (UInt ii = 0; ii < operandListCount; ++ii)
             {
-                operand->usedValue = listOperands[ii][jj];
-                operand++;
+                UInt listOperandCount = listOperandCounts[ii];
+                for (UInt jj = 0; jj < listOperandCount; ++jj)
+                {
+                    operand->usedValue = listOperands[ii][jj];
+                    operand++;
+                }
             }
         }
 
         IRInstKey key;
-        key.inst = keyInst;
+        key.inst = inst;
 
+        // Ideally we would add if not found, else return if was found instead of testing & then adding.
         IRInst* foundInst = nullptr;
         bool found = builder->sharedBuilder->globalValueNumberingMap.TryGetValue(key, foundInst);
 
-        free((void*)keyInst);
-
+        SLANG_ASSERT(endCursor == memoryArena.getCursor());
+        
         if (found)
         {
+            memoryArena.rewindToCursor(cursor);
             return foundInst;
         }
 
-        // If no instruction was found, then we need to emit it.
+        // Make the lookup instruction into proper instruction. Equivalent to
+        // IRInst* inst = createInstImpl<IRInst>(
+        //    builder,
+        //    op,
+        //    type,
+        //    0,
+        //    nullptr,
+        //    operandListCount,
+        //    listOperandCounts,
+        //    listOperands);
+
+        {
+            
+            // Okay now need to link up
+            if (type)
+            {
+                inst->typeUse.usedValue = nullptr;
+                inst->typeUse.init(inst, type);
+            }
+
+            maybeSetSourceLoc(builder, inst);
+
+            IRUse*const operands = inst->getOperands();
+            for (UInt i = 0; i < operandCount; ++i)
+            {
+                IRUse& operand = operands[i];
+                auto value = operand.usedValue;
+
+                operand.usedValue = nullptr;
+                operand.init(inst, value);
+            }
+        }
 
-        IRInst* inst = createInstImpl<IRInst>(
-            builder,
-            op,
-            type,
-            0,
-            nullptr,
-            operandListCount,
-            listOperandCounts,
-            listOperands);
         addHoistableInst(builder, inst);
 
         key.inst = inst;
diff --git a/tools/slang-test/unit-test-memory-arena.cpp b/tools/slang-test/unit-test-memory-arena.cpp
index 1b09ae9e7..e993d4f7a 100644
--- a/tools/slang-test/unit-test-memory-arena.cpp
+++ b/tools/slang-test/unit-test-memory-arena.cpp
@@ -169,6 +169,21 @@ static void memoryArenaUnitTest()
                         arena.reset();
                         blocks.Clear();
                     }
+                    else if (var == 3)
+                    {
+                        arena.rewindToCursor(nullptr);
+                        blocks.Clear();
+                    }
+                    else if (var == 4)
+                    {
+                        // Rewind to a random position
+                        int rewindIndex = randGen.nextInt32UpTo(int32_t(blocks.Count()));
+                        // rewind to this block
+                        arena.rewindToCursor(blocks[rewindIndex].m_data);
+                        // All the blocks (includign this one) and now deallocated
+                        blocks.SetSize(rewindIndex);
+
+                    }
                     else
                     {
                         size_t usedMemory = arena.calcTotalMemoryUsed();
@@ -248,6 +263,12 @@ static void memoryArenaUnitTest()
             }
         }
     }
+    {
+        // Do lots of allocations and test out rewind
+        
+
+
+    }
 }
 
-SLANG_UNIT_TEST("MemoryArena", memoryArenaUnitTest);
\ No newline at end of file
+SLANG_UNIT_TEST("MemoryArena", memoryArenaUnitTest);