From feadc830500de86c3d6ce015e33576fb5a5df4a7 Mon Sep 17 00:00:00 2001
From: Connor McLaughlin <stenzek@gmail.com>
Date: Fri, 27 Nov 2020 23:09:10 +1000
Subject: [PATCH] Common/GrowableMemoryByteStream: Add helper methods

---
 src/common/byte_stream.cpp | 61 ++++++++++++++++++++++++++++++--------
 src/common/byte_stream.h   |  7 ++++-
 2 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/src/common/byte_stream.cpp b/src/common/byte_stream.cpp
index 31f85d524..20d00ba02 100644
--- a/src/common/byte_stream.cpp
+++ b/src/common/byte_stream.cpp
@@ -700,6 +700,54 @@ GrowableMemoryByteStream::~GrowableMemoryByteStream()
     std::free(m_pPrivateMemory);
 }
 
+void GrowableMemoryByteStream::Resize(u32 new_size)
+{
+  if (new_size > m_iMemorySize)
+    ResizeMemory(new_size);
+
+  m_iSize = new_size;
+}
+
+void GrowableMemoryByteStream::ResizeMemory(u32 new_size)
+{
+  if (new_size == m_iMemorySize)
+    return;
+
+  if (m_pPrivateMemory == nullptr)
+  {
+    m_pPrivateMemory = (u8*)std::malloc(new_size);
+    std::memcpy(m_pPrivateMemory, m_pMemory, m_iSize);
+    m_pMemory = m_pPrivateMemory;
+    m_iMemorySize = new_size;
+  }
+  else
+  {
+    m_pPrivateMemory = m_pMemory = (u8*)std::realloc(m_pPrivateMemory, new_size);
+    m_iMemorySize = new_size;
+  }
+}
+
+void GrowableMemoryByteStream::EnsureSpace(u32 space)
+{
+  if ((m_iSize + space) >= m_iMemorySize)
+    return;
+
+  Grow((m_iSize + space) - m_iMemorySize);
+}
+
+void GrowableMemoryByteStream::ShrinkToFit()
+{
+  if (!m_pPrivateMemory || m_iSize == m_iMemorySize)
+    return;
+
+  u8* new_ptr = static_cast<u8*>(std::realloc(m_pPrivateMemory, m_iSize));
+  if (new_ptr)
+  {
+    m_pPrivateMemory = new_ptr;
+    m_iMemorySize = m_iSize;
+  }
+}
+
 bool GrowableMemoryByteStream::ReadByte(u8* pDestByte)
 {
   if (m_iPosition < m_iSize)
@@ -819,18 +867,7 @@ bool GrowableMemoryByteStream::Discard()
 void GrowableMemoryByteStream::Grow(u32 MinimumGrowth)
 {
   u32 NewSize = std::max(m_iMemorySize + MinimumGrowth, m_iMemorySize * 2);
-  if (m_pPrivateMemory == nullptr)
-  {
-    m_pPrivateMemory = (u8*)std::malloc(NewSize);
-    std::memcpy(m_pPrivateMemory, m_pMemory, m_iSize);
-    m_pMemory = m_pPrivateMemory;
-    m_iMemorySize = NewSize;
-  }
-  else
-  {
-    m_pPrivateMemory = m_pMemory = (u8*)std::realloc(m_pPrivateMemory, NewSize);
-    m_iMemorySize = NewSize;
-  }
+  ResizeMemory(NewSize);
 }
 
 #if defined(_MSC_VER)
diff --git a/src/common/byte_stream.h b/src/common/byte_stream.h
index 14ce021e6..a9096a12f 100644
--- a/src/common/byte_stream.h
+++ b/src/common/byte_stream.h
@@ -168,7 +168,12 @@ public:
   virtual ~GrowableMemoryByteStream();
 
   u8* GetMemoryPointer() const { return m_pMemory; }
-  u32 GetMemorySize() const { return m_iSize; }
+  u32 GetMemorySize() const { return m_iMemorySize; }
+
+  void Resize(u32 new_size);
+  void ResizeMemory(u32 new_size);
+  void EnsureSpace(u32 space);
+  void ShrinkToFit();
 
   virtual bool ReadByte(u8* pDestByte) override;
   virtual u32 Read(void* pDestination, u32 ByteCount) override;