Riff Container Stream Writing (#1116)

* * Added option to get random bytes from RandomGenerator
* Added ability to allocate only in current block on MemoryArena
* Allowed RiffContainer to not allocate new Data blocks, if can just extend the Data it has (because it's at the end of current block and there is space for the new data).
* Added coverage for change on Riff unit test

* Add test coverage for allocations over multiple Data blocks.

* Improve comment on riff unit tests.

6 files changed, 302 insertions, 6 deletions

diff --git a/source/core/slang-memory-arena.h b/source/core/slang-memory-arena.h
index d346197f0..f63a48413 100644
--- a/source/core/slang-memory-arena.h
+++ b/source/core/slang-memory-arena.h
@@ -84,6 +84,11 @@ public:
         @return The allocation (or nullptr if unable to allocate).  */
     void* allocateUnaligned(size_t sizeInBytes);
 
+        /** Allocate some aligned memory of at least size bytes, without alignment, and only from current block.
+        @param sizeInBytes Size of allocation wanted. 
+        @return The allocation (or nullptr if unable to allocate in current block).  */
+    void* allocateCurrentUnaligned(size_t sizeInBytes);
+
         /** Allocates a null terminated string.
 
         NOTE, it is not possible to rewind to a zero length string allocation (because such a strings memory is not held on the arena)
@@ -127,6 +132,9 @@ public:
         /// Gets the block alignment that is passed at initialization otherwise 0 an invalid block alignment.
     size_t getBlockAlignment() const { return m_blockAlignment; }
 
+        /// Get the default block payload size
+    size_t getBlockPayloadSize() const { return m_blockPayloadSize; }
+
         /// 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
@@ -241,6 +249,23 @@ SLANG_FORCE_INLINE void* MemoryArena::allocateUnaligned(size_t sizeInBytes)
 }
 
 // --------------------------------------------------------------------------
+SLANG_FORCE_INLINE void* MemoryArena::allocateCurrentUnaligned(size_t sizeInBytes)
+{
+    // Align with the minimum alignment
+    uint8_t* mem = m_current;
+    uint8_t* end = mem + sizeInBytes;
+    if (end <= m_end)
+    {
+        m_current = end;
+        return mem;
+    }
+    else
+    {
+        return nullptr;
+    }
+}
+
+// --------------------------------------------------------------------------
 SLANG_FORCE_INLINE void* MemoryArena::allocate(size_t sizeInBytes)
 {
     assert(sizeInBytes > 0);
diff --git a/source/core/slang-random-generator.cpp b/source/core/slang-random-generator.cpp
index ce43067aa..ec06336f1 100644
--- a/source/core/slang-random-generator.cpp
+++ b/source/core/slang-random-generator.cpp
@@ -71,6 +71,57 @@ int64_t RandomGenerator::nextInt64InRange(int64_t min, int64_t max)
     return (nextPositiveInt64() % diff) + min;
 }
 
+static uint8_t* _nextData(RandomGenerator* rand, uint8_t* out, size_t size)
+{
+    if (size)
+    {
+        SLANG_ASSERT(size <= 4);
+        uint32_t v = uint32_t(rand->nextInt32());
+        uint8_t* dst = (uint8_t*)out;
+        for (size_t i = 0; i < size; ++i)
+        {
+            dst[i] = uint8_t(v);
+            v >>= 8;
+        }
+    }
+    return out + size;
+}
+
+void RandomGenerator::nextData(void* out, size_t size)
+{
+    uint8_t* dst = (uint8_t*)out;
+    uint8_t*const end = dst + size;
+
+    // For short runs just output
+    if (size <= 4)
+    {
+        _nextData(this, dst, size);
+        return;
+    }
+    
+    {
+        const size_t preAlign = size_t(((size_t(dst) + 3) & ~size_t(3)) - size_t(dst));
+        dst = _nextData(this, dst, preAlign);
+    }
+
+    // Check invariants
+    SLANG_ASSERT((size_t(dst) & 3) == 0 && end >= dst);
+
+    {
+        const size_t middleCount = size_t(end - dst) >> 2;
+        if (middleCount)
+        {
+            nextInt32s((int32_t*)dst, middleCount);
+            dst += middleCount * sizeof(int32_t);
+        }
+    }
+
+    // Check invariants
+    SLANG_ASSERT((size_t(dst) & 3) == 0 && end >= dst);
+
+    _nextData(this, dst, size_t(end - dst));
+}
+
 /* static */RandomGenerator* RandomGenerator::create(int32_t seed)
 {
     return new DefaultRandomGenerator(seed);
@@ -155,7 +206,34 @@ int32_t Mt19937RandomGenerator::nextInt32()
     return int32_t(y);
 }
 
+void Mt19937RandomGenerator::nextInt32s(int32_t* dst, size_t count)
+{
+    while (count)
+    {
+        if (m_index >= kNumEntries)
+        {
+            _generate();
+        }
+
+        const size_t remaining = kNumEntries - m_index;
+        const size_t run = (count < remaining) ? count : remaining;
+
+        const uint32_t* src = m_mt + m_index;
+        for (size_t i = 0; i < run; i++)
+        {
+            uint32_t y = src[i];
+            y = y ^ (y >> 11);
+            y = y ^ ((y << 7) & uint32_t(0x9d2c5680u));
+            y = y ^ ((y << 15) & uint32_t(0xefc6000u));
+            y = y ^ (y >> 18);
 
+            dst[i] = int32_t(y);
+        }
 
+        m_index += int(run);
+        dst += run;
+        count -= run;
+    }
+}
 
 } // namespace Slang
diff --git a/source/core/slang-random-generator.h b/source/core/slang-random-generator.h
index 57f0e8630..392b4cb4c 100644
--- a/source/core/slang-random-generator.h
+++ b/source/core/slang-random-generator.h
@@ -30,6 +30,9 @@ class RandomGenerator: public RefObject
         /// Get the next bool
     virtual bool nextBool();
 
+        /// Get multiple int32s 
+    virtual void nextInt32s(int32_t* dst, size_t count) = 0;
+
         /// Next uint32_t
     uint32_t nextUInt32() { return uint32_t(nextInt32()); }
 
@@ -53,6 +56,10 @@ class RandomGenerator: public RefObject
         /// Returns value from min up to BUT NOT INCLUDING max
     int64_t nextInt64InRange(int64_t min, int64_t max);
 
+        /// Fill with random data.
+        /// NOTE! Output is only identical bytes if generator in same state *and* size_t(dst) & 3 is the same on calls. 
+    void nextData(void* dst, size_t size);
+
         /// Create a RandomGenerator with specified seed using default generator type
     static RandomGenerator* create(int32_t seed);
 };
@@ -73,7 +80,8 @@ class Mt19937RandomGenerator: public RandomGenerator
     Mt19937RandomGenerator* clone() SLANG_OVERRIDE { return new ThisType(*this); }
     void reset(int32_t seed) SLANG_OVERRIDE;
     int32_t nextInt32() SLANG_OVERRIDE;
-    
+    void nextInt32s(int32_t* dst, size_t count) SLANG_OVERRIDE;
+
         /// Ctor
     Mt19937RandomGenerator();
     Mt19937RandomGenerator(const ThisType& rhs);
diff --git a/source/core/slang-riff.cpp b/source/core/slang-riff.cpp
index 1de51d840..da547319b 100644
--- a/source/core/slang-riff.cpp
+++ b/source/core/slang-riff.cpp
@@ -665,6 +665,47 @@ size_t RiffContainer::DataChunk::calcPayloadSize() const
     return size;
 }
 
+void RiffContainer::DataChunk::getPayload(void* inDst) const
+{
+    uint8_t* dst = (uint8_t*)inDst;
+
+    Data* data = m_dataList;
+    while (data)
+    {
+        const size_t size = data->getSize();
+        ::memcpy(dst, data->getPayload(), size);
+
+        dst += size;
+        data = data->m_next;
+    }
+}
+
+bool RiffContainer::DataChunk::isEqual(const void* inData, size_t count) const
+{
+    const uint8_t* src = (const uint8_t*)inData;
+
+    Data* data = m_dataList;
+    while (data)
+    {
+        const size_t size = data->getSize();
+        // Can't have more content than remaining
+        // Contents must match
+        if (size > count || ::memcmp(src, data->getPayload(), size) != 0)
+        {
+            return false;
+        }
+
+        src += size;
+        count -= size;
+
+        // Next data block
+        data = data->m_next;
+    }
+
+    // If match must be at the end
+    return count == 0;
+}
+
 // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! RiffContainer !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 
 RiffContainer::RiffContainer() :
@@ -862,8 +903,63 @@ RiffContainer::Data* RiffContainer::addData()
     return data;
 }
 
+RiffContainer::Data* RiffContainer::makeSingleData(DataChunk* dataChunk)
+{
+    // There is no data
+    if (dataChunk->m_dataList == nullptr)
+    {
+        return nullptr;
+    }
+
+    if (dataChunk->m_dataList->m_next == nullptr)
+    {
+        return dataChunk->m_dataList;
+    }
+
+    {
+        Data* data = dataChunk->m_dataList;
+        
+        // Okay lets combine all into one block
+        const size_t payloadSize = dataChunk->calcPayloadSize();
+
+        void* dst = m_arena.allocate(payloadSize);
+        dataChunk->getPayload(dst);
+
+        // Remove other datas
+        data->m_next = nullptr;
+        // Make this the end
+        dataChunk->m_endData = data;
+
+        // Point to the block with all of the data
+        data->m_ownership = Ownership::Arena;
+        data->m_payload = dst;
+        data->m_size = payloadSize;
+
+        return data;
+    }
+}
+
 void RiffContainer::write(const void* inData, size_t size)
 {
+    // We must be in a chunk
+    SLANG_ASSERT(m_dataChunk);
+    // Get the last data chunk
+    Data* endData = m_dataChunk->m_endData;
+    if (endData)
+    {
+        uint8_t* end = ((uint8_t*)endData->m_payload) + endData->m_size;
+        // See if can just add to end of current data
+        if ( end == m_arena.getCursor() && m_arena.allocateCurrentUnaligned(size))
+        {
+            ::memcpy(end, inData, size);
+            endData->m_size += size;
+
+            // Add current chunks data
+            m_dataChunk->m_payloadSize += size;
+            return;
+        }
+    }
+
     auto data = addData();
     setPayload(data, inData, size);
 }
diff --git a/source/core/slang-riff.h b/source/core/slang-riff.h
index 61e6eed1e..1622f60f4 100644
--- a/source/core/slang-riff.h
+++ b/source/core/slang-riff.h
@@ -284,6 +284,12 @@ public:
             /// Calculate the payload size
         size_t calcPayloadSize() const;
 
+            /// Copy the payload to dst. Dst must be at least the payload size. 
+        void getPayload(void* dst) const;
+
+            /// True if payloads contents is equal to data
+        bool isEqual(const void* data, size_t count) const;
+
             /// Get single data payload.
         Data* getSingleData() const;
 
@@ -347,12 +353,21 @@ public:
         /// Get the root
     ListChunk* getRoot() const { return m_rootList; }
 
+        /// Get the current chunk
+    Chunk* getCurrentChunk() { return m_dataChunk ? static_cast<Chunk*>(m_dataChunk) : static_cast<Chunk*>(m_listChunk); }
+
         /// Reset the container
     void reset();
 
         /// true if has a root container, and nothing remains open
     bool isFullyConstructed() { return m_rootList && m_listChunk == nullptr && m_dataChunk == nullptr; }
 
+        /// Makes a data chunk contain a single contiguous data block
+    Data* makeSingleData(DataChunk* dataChunk);
+
+        /// Get the memory arena that is backing the storage of data
+    MemoryArena& getMemoryArena() { return m_arena; }
+
         /// The if the list and sublists appear correct
     static bool isChunkOk(Chunk* chunk);
 
diff --git a/tools/slang-test/unit-test-riff.cpp b/tools/slang-test/unit-test-riff.cpp
index 58e9a1bb7..0fb81113a 100644
--- a/tools/slang-test/unit-test-riff.cpp
+++ b/tools/slang-test/unit-test-riff.cpp
@@ -2,24 +2,50 @@
 
 #include "../../source/core/slang-riff.h"
 
+#include "../../source/core/slang-random-generator.h"
+
 #include "test-context.h"
 
 using namespace Slang;
 
+static void _writeRandom(RandomGenerator* rand, size_t maxSize, RiffContainer& ioContainer, List<uint8_t>& ioData)
+{
+    while (true)
+    {
+        const Index oldCount = ioData.getCount();
+
+        const size_t allocSize = size_t(rand->nextInt32InRange(1, 50));
+
+        if (allocSize + oldCount > maxSize)
+        {
+            break;
+        }
+
+        ioData.setCount(oldCount + Index(allocSize));
+        rand->nextData(ioData.getBuffer() + oldCount, allocSize);
+
+        // Write
+        ioContainer.write(ioData.getBuffer() + oldCount, allocSize);
+    }
+
+    // Should be a single block with same data as the List
+    RiffContainer::DataChunk* dataChunk = as<RiffContainer::DataChunk>(ioContainer.getCurrentChunk());
+    SLANG_ASSERT(dataChunk);
+}
+
 static void riffUnitTest()
 {
+    typedef RiffContainer::ScopeChunk ScopeChunk;
+    typedef RiffContainer::Chunk::Kind Kind;
+
     const FourCC markThings = SLANG_FOUR_CC('T', 'H', 'I', 'N');
     const FourCC markData = SLANG_FOUR_CC('D', 'A', 'T', 'A');
 
     {
-        typedef RiffContainer::ScopeChunk ScopeChunk;
-        typedef RiffContainer::Chunk::Kind Kind;
-
         RiffContainer container;
 
         {
-            ScopeChunk scopeContainer(&container, Kind::List, markThings);
-
+            ScopeChunk scopeContainer(&container, Kind::List, markThings);        
             {
                 ScopeChunk scopeChunk(&container, Kind::Data, markData);
 
@@ -85,6 +111,54 @@ static void riffUnitTest()
 
     }
 
+    // Test writing as a stream only allocates a single data block (as long as there is enough space).
+    {
+        RiffContainer container;
+
+        ScopeChunk scopeChunk(&container, Kind::List, markData);
+        {
+            ScopeChunk scopeChunk(&container, Kind::Data, markData);
+            RefPtr<RandomGenerator> rand = RandomGenerator::create(0x345234);
+
+            List<uint8_t> data;
+            _writeRandom(rand, container.getMemoryArena().getBlockPayloadSize() / 2, container, data); 
+
+            // Should be a single block with same data as the List
+            RiffContainer::DataChunk* dataChunk = as<RiffContainer::DataChunk>(container.getCurrentChunk());
+            SLANG_ASSERT(dataChunk);
+
+            // It should be a single block
+            SLANG_CHECK(dataChunk->getSingleData() != nullptr);
+
+            SLANG_CHECK(dataChunk->isEqual(data.getBuffer(), data.getCount()));
+
+        }
+    } 
+
+    // Test writing across multiple data blocks
+    {
+        RefPtr<RandomGenerator> rand = RandomGenerator::create(0x345234);
+
+        for (Int i = 0 ; i < 100; ++i)
+        {
+            RiffContainer container;
+
+            const size_t maxSize = rand->nextInt32InRange(1, int32_t(container.getMemoryArena().getBlockPayloadSize() * 3));
+            
+            ScopeChunk scopeChunk(&container, Kind::List, markData);
+            {
+                ScopeChunk scopeChunk(&container, Kind::Data, markData);
+            
+                List<uint8_t> data;
+                _writeRandom(rand, maxSize, container, data);
+
+                // Should be a single block with same data as the List
+                RiffContainer::DataChunk* dataChunk = as<RiffContainer::DataChunk>(container.getCurrentChunk());
+                SLANG_CHECK(dataChunk && dataChunk->isEqual(data.getBuffer(), data.getCount()));
+            }
+        }
+    }
+
 #if 0
     {
         RiffContainer container;