From f708e5c0b28820820105b87536f074ecf533f63a Mon Sep 17 00:00:00 2001
From: zilmar <zilmar@pj64-emu.com>
Date: Thu, 17 Oct 2024 15:05:48 +1030
Subject: [PATCH] Core: Check recompiler memory based on the function size

---
 .../N64System/Recompiler/CodeBlock.cpp          | 15 +++++++++------
 .../N64System/Recompiler/CodeBlock.h            |  3 ++-
 .../N64System/Recompiler/Recompiler.cpp         | 12 +++++++-----
 .../N64System/Recompiler/RecompilerMemory.cpp   | 17 +++++++++++++----
 .../N64System/Recompiler/RecompilerMemory.h     |  7 ++++---
 5 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/Source/Project64-core/N64System/Recompiler/CodeBlock.cpp b/Source/Project64-core/N64System/Recompiler/CodeBlock.cpp
index 9aee8a5eb..615d63afb 100644
--- a/Source/Project64-core/N64System/Recompiler/CodeBlock.cpp
+++ b/Source/Project64-core/N64System/Recompiler/CodeBlock.cpp
@@ -892,25 +892,28 @@ bool CCodeBlock::Compile()
     return true;
 }
 
-uint32_t CCodeBlock::Finilize(uint8_t * CompiledLocation)
+uint32_t CCodeBlock::Finilize(CRecompMemory & RecompMem)
 {
+    m_CompiledLocation = RecompMem.RecompPos();
+    m_CodeHolder.relocateToBase((uint64_t)m_CompiledLocation);
     if (CDebugSettings::bRecordRecompilerAsm())
     {
         std::string CodeLog = m_CodeLog;
         m_CodeLog.clear();
         Log("====== Code block ======");
-        Log("Native entry point: %X", CompiledLocation);
+        Log("Native entry point: %X", m_CompiledLocation);
         Log("Start of block: %X", VAddrEnter());
         Log("Number of sections: %d", NoOfSections());
         Log("====== Recompiled code ======");
         m_CodeLog += CodeLog;
     }
-
-    m_CompiledLocation = CompiledLocation;
-    m_CodeHolder.relocateToBase((uint64_t)m_CompiledLocation);
     size_t codeSize = m_CodeHolder.codeSize();
+    if (!RecompMem.CheckRecompMem(codeSize))
+    {
+        return 0;
+    }
     m_CodeHolder.copyFlattenedData(m_CompiledLocation, codeSize, asmjit::CopySectionFlags::kPadSectionBuffer);
-    *m_Recompiler.RecompPos() += codeSize;
+    m_Recompiler.RecompPos() += codeSize;
 
 #if defined(ANDROID) && (defined(__arm__) || defined(_M_ARM))
     __clear_cache((uint8_t *)((uint32_t)m_CompiledLocation & ~1), m_CompiledLocation + codeSize);
diff --git a/Source/Project64-core/N64System/Recompiler/CodeBlock.h b/Source/Project64-core/N64System/Recompiler/CodeBlock.h
index eeced9469..d10b88968 100644
--- a/Source/Project64-core/N64System/Recompiler/CodeBlock.h
+++ b/Source/Project64-core/N64System/Recompiler/CodeBlock.h
@@ -9,6 +9,7 @@
 
 class CMipsMemoryVM;
 class CRecompiler;
+class CRecompMemory;
 
 class CCodeBlock :
     public asmjit::ErrorHandler
@@ -18,7 +19,7 @@ public:
     ~CCodeBlock();
 
     bool Compile();
-    uint32_t Finilize(uint8_t * CompiledLocation);
+    uint32_t Finilize(CRecompMemory & RecompMem);
 
     asmjit::CodeHolder & CodeHolder(void)
     {
diff --git a/Source/Project64-core/N64System/Recompiler/Recompiler.cpp b/Source/Project64-core/N64System/Recompiler/Recompiler.cpp
index 2e44121a1..622829fc3 100644
--- a/Source/Project64-core/N64System/Recompiler/Recompiler.cpp
+++ b/Source/Project64-core/N64System/Recompiler/Recompiler.cpp
@@ -356,7 +356,6 @@ CCompiledFunc * CRecompiler::CompileCode()
         }
     }
 
-    CheckRecompMem();
     WriteTrace(TraceRecompiler, TraceDebug, "Compile Block-Start: Program Counter: %016llX pAddr: %X", PROGRAM_COUNTER, pAddr);
 
     CCodeBlock CodeBlock(m_System, (uint32_t)PROGRAM_COUNTER);
@@ -364,16 +363,19 @@ CCompiledFunc * CRecompiler::CompileCode()
     {
         return nullptr;
     }
+    uint32_t CodeLen = CodeBlock.Finilize(*this);
+    if (CodeLen == 0)
+    {
+        return nullptr;
+    }
+    RecompPos() += CodeLen;
+    LogCodeBlock(CodeBlock);
 
     if (bShowRecompMemSize())
     {
         ShowMemUsed();
     }
 
-    uint32_t CodeLen = CodeBlock.Finilize(*RecompPos());
-    *RecompPos() += CodeLen;
-    LogCodeBlock(CodeBlock);
-
     if (bSMM_StoreInstruc())
     {
         m_MMU.ClearMemoryWriteMap(CodeBlock.VAddrEnter() & ~0xFFF, 0xFFF);
diff --git a/Source/Project64-core/N64System/Recompiler/RecompilerMemory.cpp b/Source/Project64-core/N64System/Recompiler/RecompilerMemory.cpp
index 05b28aaf1..59dbbd46b 100644
--- a/Source/Project64-core/N64System/Recompiler/RecompilerMemory.cpp
+++ b/Source/Project64-core/N64System/Recompiler/RecompilerMemory.cpp
@@ -49,17 +49,25 @@ bool CRecompMemory::AllocateMemory()
     return true;
 }
 
-void CRecompMemory::CheckRecompMem()
+bool CRecompMemory::CheckRecompMem(uint32_t BlockSize)
 {
     uint32_t Size = (uint32_t)((uint8_t *)m_RecompPos - (uint8_t *)m_RecompCode);
-    if ((Size + 0x50000) < m_RecompSize)
+    if (BlockSize < 0x50000)
     {
-        return;
+        BlockSize = 0x50000;
+    }
+    if ((Size + BlockSize) < m_RecompSize)
+    {
+        return true;
     }
     if (m_RecompSize == MaxCompileBufferSize)
     {
         g_Recompiler->ResetRecompCode(true);
-        return;
+        return false;
+    }
+    if (BlockSize > IncreaseCompileBufferSize)
+    {
+        g_Notify->BreakPoint(__FILE__, __LINE__);
     }
     void * MemAddr = CommitMemory(m_RecompCode + m_RecompSize, IncreaseCompileBufferSize, MEM_EXECUTE_READWRITE);
     if (MemAddr == nullptr)
@@ -68,6 +76,7 @@ void CRecompMemory::CheckRecompMem()
         g_Notify->FatalError(MSG_MEM_ALLOC_ERROR);
     }
     m_RecompSize += IncreaseCompileBufferSize;
+    return true;
 }
 
 void CRecompMemory::Reset()
diff --git a/Source/Project64-core/N64System/Recompiler/RecompilerMemory.h b/Source/Project64-core/N64System/Recompiler/RecompilerMemory.h
index 029dda0fd..7e76af9c6 100644
--- a/Source/Project64-core/N64System/Recompiler/RecompilerMemory.h
+++ b/Source/Project64-core/N64System/Recompiler/RecompilerMemory.h
@@ -8,16 +8,17 @@ protected:
     ~CRecompMemory();
 
     bool AllocateMemory();
-    void CheckRecompMem();
     void Reset();
     void ShowMemUsed();
 
 public:
-    uint8_t ** RecompPos()
+    uint8_t *& RecompPos()
     {
-        return &m_RecompPos;
+        return m_RecompPos;
     }
 
+    bool CheckRecompMem(uint32_t BlockSize);
+
 private:
     CRecompMemory(const CRecompMemory &);
     CRecompMemory & operator=(const CRecompMemory &);