From 39f19e1a8f05538dc67bd9bcc7c999beeefcb64d Mon Sep 17 00:00:00 2001
From: kobalicek <kobalicek.petr@gmail.com>
Date: Sat, 2 May 2020 16:54:11 +0200
Subject: [PATCH] [Bug] Fixed #285 and #241

---
 src/asmjit/core/compiler.cpp | 30 ++++++++++++++++++++++++++++++
 src/asmjit/core/compiler.h   |  8 ++++++++
 src/asmjit/core/rapass.cpp   |  7 ++++---
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/asmjit/core/compiler.cpp b/src/asmjit/core/compiler.cpp
index 5b7e0db..307c663 100644
--- a/src/asmjit/core/compiler.cpp
+++ b/src/asmjit/core/compiler.cpp
@@ -12,6 +12,7 @@
 #include "../core/cpuinfo.h"
 #include "../core/logging.h"
 #include "../core/rapass_p.h"
+#include "../core/rastack_p.h"
 #include "../core/support.h"
 #include "../core/type.h"
 
@@ -445,6 +446,35 @@ Error BaseCompiler::_newStack(BaseMem& out, uint32_t size, uint32_t alignment, c
   return kErrorOk;
 }
 
+Error BaseCompiler::setStackSize(uint32_t virtId, uint32_t newSize, uint32_t newAlignment) noexcept {
+  if (!isVirtIdValid(virtId))
+    return DebugUtils::errored(kErrorInvalidVirtId);
+
+  if (newAlignment && !Support::isPowerOf2(newAlignment))
+    return reportError(DebugUtils::errored(kErrorInvalidArgument));
+
+  if (newAlignment > 64)
+    newAlignment = 64;
+
+  VirtReg* vReg = virtRegById(virtId);
+  if (newSize)
+    vReg->_virtSize = newSize;
+
+  if (newAlignment)
+    vReg->_alignment = uint8_t(newAlignment);
+
+  // This is required if the RAPass is already running. There is a chance that
+  // a stack-slot has been already allocated and in that case it has to be
+  // updated as well, otherwise we would allocate wrong amount of memory.
+  RAWorkReg* workReg = vReg->_workReg;
+  if (workReg && workReg->_stackSlot) {
+    workReg->_stackSlot->_size = vReg->virtSize();
+    workReg->_stackSlot->_alignment = vReg->alignment();
+  }
+
+  return kErrorOk;
+}
+
 Error BaseCompiler::_newConst(BaseMem& out, uint32_t scope, const void* data, size_t size) {
   ConstPoolNode** pPool;
   if (scope == ConstPool::kScopeLocal)
diff --git a/src/asmjit/core/compiler.h b/src/asmjit/core/compiler.h
index 0447345..b4dd4a8 100644
--- a/src/asmjit/core/compiler.h
+++ b/src/asmjit/core/compiler.h
@@ -282,6 +282,14 @@ public:
 
   ASMJIT_API Error _newStack(BaseMem& out, uint32_t size, uint32_t alignment, const char* name = nullptr);
 
+  //! Updates the stack size of a stack created by `_newStack()` by its `virtId`.
+  ASMJIT_API Error setStackSize(uint32_t virtId, uint32_t newSize, uint32_t newAlignment = 0) noexcept;
+
+  //! Updates the stack size of a stack created by `_newStack()`.
+  inline Error setStackSize(const BaseMem& mem, uint32_t newSize, uint32_t newAlignment = 0) noexcept {
+    return setStackSize(mem.id(), newSize, newAlignment);
+  }
+
   //! \}
 
   //! \name Constants
diff --git a/src/asmjit/core/rapass.cpp b/src/asmjit/core/rapass.cpp
index 0c0ea73..129a1d5 100644
--- a/src/asmjit/core/rapass.cpp
+++ b/src/asmjit/core/rapass.cpp
@@ -1429,11 +1429,12 @@ Error RAPass::useTemporaryMem(BaseMem& out, uint32_t size, uint32_t alignment) n
   }
   else {
     ASMJIT_ASSERT(_temporaryMem.as<BaseMem>().isRegHome());
-    uint32_t virtId = _temporaryMem.as<BaseMem>().baseId();
 
+    uint32_t virtId = _temporaryMem.as<BaseMem>().baseId();
     VirtReg* virtReg = cc()->virtRegById(virtId);
-    virtReg->_virtSize = Support::max(virtReg->virtSize(), size);
-    virtReg->_alignment = uint8_t(Support::max(virtReg->alignment(), alignment));
+
+    cc()->setStackSize(virtId, Support::max(virtReg->virtSize(), size),
+                               Support::max(virtReg->alignment(), alignment));
   }
 
   out = _temporaryMem.as<BaseMem>();