JITs: Add EraseSingleBlock Function

Source/Core/Core/PowerPC/CachedInterpreter/CachedInterpreter.cpp

@@ -402,6 +402,12 @@ bool CachedInterpreter::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
   return true;
 }
 
+void CachedInterpreter::EraseSingleBlock(const JitBlock& block)
+{
+  m_block_cache.EraseSingleBlock(block);
+  FreeRanges();
+}
+
 void CachedInterpreter::ClearCache()
 {
   m_block_cache.Clear();

Source/Core/Core/PowerPC/CachedInterpreter/CachedInterpreter.h

@@ -46,6 +46,8 @@ public:
   void Jit(u32 address, bool clear_cache_and_retry_on_failure);
   bool DoJit(u32 address, JitBlock* b, u32 nextPC);
 
+  void EraseSingleBlock(const JitBlock& block) override;
+
   static std::size_t Disassemble(const JitBlock& block, std::ostream& stream);
 
   JitBaseBlockCache* GetBlockCache() override { return &m_block_cache; }

Source/Core/Core/PowerPC/Jit64/Jit.cpp

@@ -1202,6 +1202,12 @@ bool Jit64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
   return true;
 }
 
+void Jit64::EraseSingleBlock(const JitBlock& block)
+{
+  blocks.EraseSingleBlock(block);
+  FreeRanges();
+}
+
 BitSet8 Jit64::ComputeStaticGQRs(const PPCAnalyst::CodeBlock& cb) const
 {
   return cb.m_gqr_used & ~cb.m_gqr_modified;

Source/Core/Core/PowerPC/Jit64/Jit.h

@@ -65,6 +65,8 @@ public:
   void Jit(u32 em_address, bool clear_cache_and_retry_on_failure);
   bool DoJit(u32 em_address, JitBlock* b, u32 nextPC);
 
+  void EraseSingleBlock(const JitBlock& block) override;
+
   // Finds a free memory region and sets the near and far code emitters to point at that region.
   // Returns false if no free memory region can be found for either of the two.
   bool SetEmitterStateToFreeCodeRegion();

Source/Core/Core/PowerPC/JitArm64/Jit.cpp

@@ -1034,6 +1034,12 @@ void JitArm64::Jit(u32 em_address, bool clear_cache_and_retry_on_failure)
   exit(-1);
 }
 
+void JitArm64::EraseSingleBlock(const JitBlock& block)
+{
+  blocks.EraseSingleBlock(block);
+  FreeRanges();
+}
+
 std::optional<size_t> JitArm64::SetEmitterStateToFreeCodeRegion()
 {
   // Find some large free memory blocks and set code emitters to point at them. If we can't find

Source/Core/Core/PowerPC/JitArm64/Jit.h

@@ -48,6 +48,8 @@ public:
   void Jit(u32 em_address) override;
   void Jit(u32 em_address, bool clear_cache_and_retry_on_failure);
 
+  void EraseSingleBlock(const JitBlock& block) override;
+
   const char* GetName() const override { return "JITARM64"; }
 
   // OPCODES

Source/Core/Core/PowerPC/JitCommon/JitBase.h

@@ -194,6 +194,8 @@ public:
 
   virtual void Jit(u32 em_address) = 0;
 
+  virtual void EraseSingleBlock(const JitBlock& block) = 0;
+
   virtual const CommonAsmRoutinesBase* GetAsmRoutines() = 0;
 
   virtual bool HandleFault(uintptr_t access_address, SContext* ctx) = 0;

Source/Core/Core/PowerPC/JitCommon/JitCache.cpp

@@ -155,11 +155,10 @@ void JitBaseBlockCache::FinalizeBlock(JitBlock& block, bool block_link,
 
   block.physical_addresses = physical_addresses;
 
-  u32 range_mask = ~(BLOCK_RANGE_MAP_ELEMENTS - 1);
+  for (u32 addr : block.physical_addresses)
-  for (u32 addr : physical_addresses)
   {
     valid_block.Set(addr / 32);
-    block_range_map[addr & range_mask].insert(&block);
+    block_range_map[addr & BLOCK_RANGE_MAP_MASK].insert(&block);
   }
 
   if (block_link)
@@ -333,8 +332,7 @@ void JitBaseBlockCache::InvalidateICacheInternal(u32 physical_address, u32 addre
 void JitBaseBlockCache::ErasePhysicalRange(u32 address, u32 length)
 {
   // Iterate over all macro blocks which overlap the given range.
-  u32 range_mask = ~(BLOCK_RANGE_MAP_ELEMENTS - 1);
+  auto start = block_range_map.lower_bound(address & BLOCK_RANGE_MAP_MASK);
-  auto start = block_range_map.lower_bound(address & range_mask);
   auto end = block_range_map.lower_bound(address + length);
   while (start != end)
   {
@@ -348,8 +346,8 @@ void JitBaseBlockCache::ErasePhysicalRange(u32 address, u32 length)
         // If the block overlaps, also remove all other occupied slots in the other macro blocks.
         // This will leak empty macro blocks, but they may be reused or cleared later on.
         for (u32 addr : block->physical_addresses)
-          if ((addr & range_mask) != start->first)
+          if ((addr & BLOCK_RANGE_MAP_MASK) != start->first)
-            block_range_map[addr & range_mask].erase(block);
+            block_range_map[addr & BLOCK_RANGE_MAP_MASK].erase(block);
 
         // And remove the block.
         DestroyBlock(*block);
@@ -379,6 +377,23 @@ void JitBaseBlockCache::ErasePhysicalRange(u32 address, u32 length)
   }
 }
 
+void JitBaseBlockCache::EraseSingleBlock(const JitBlock& block)
+{
+  const auto equal_range = block_map.equal_range(block.physicalAddress);
+  const auto block_map_iter = std::ranges::find(equal_range.first, equal_range.second, &block,
+                                                [](const auto& kv) { return &kv.second; });
+  if (block_map_iter == equal_range.second) [[unlikely]]
+    return;
+
+  JitBlock& mutable_block = block_map_iter->second;
+
+  for (const u32 addr : mutable_block.physical_addresses)
+    block_range_map[addr & BLOCK_RANGE_MAP_MASK].erase(&mutable_block);
+
+  DestroyBlock(mutable_block);
+  block_map.erase(block_map_iter);  // The original JitBlock reference is now dangling.
+}
+
 u32* JitBaseBlockCache::GetBlockBitSet() const
 {
   return valid_block.m_valid_block.get();

Source/Core/Core/PowerPC/JitCommon/JitCache.h

@@ -180,6 +180,7 @@ public:
   void InvalidateICache(u32 address, u32 length, bool forced);
   void InvalidateICacheLine(u32 address);
   void ErasePhysicalRange(u32 address, u32 length);
+  void EraseSingleBlock(const JitBlock& block);
 
   u32* GetBlockBitSet() const;
 
@@ -213,7 +214,7 @@ private:
   // Range of overlapping code indexed by a masked physical address.
   // This is used for invalidation of memory regions. The range is grouped
   // in macro blocks of each 0x100 bytes.
-  static constexpr u32 BLOCK_RANGE_MAP_ELEMENTS = 0x100;
+  static constexpr u32 BLOCK_RANGE_MAP_MASK = ~(0x100 - 1);
   std::map<u32, std::unordered_set<JitBlock*>> block_range_map;
 
   // This bitsets shows which cachelines overlap with any blocks.

Source/Core/Core/PowerPC/JitInterface.cpp

@@ -263,6 +263,12 @@ void JitInterface::ClearSafe()
     m_jit->GetBlockCache()->Clear();
 }
 
+void JitInterface::EraseSingleBlock(const JitBlock& block)
+{
+  if (m_jit)
+    m_jit->EraseSingleBlock(block);
+}
+
 void JitInterface::InvalidateICache(u32 address, u32 size, bool forced)
 {
   if (m_jit)

Source/Core/Core/PowerPC/JitInterface.h

@@ -16,6 +16,7 @@
 class CPUCoreBase;
 class PointerWrap;
 class JitBase;
+struct JitBlock;
 
 namespace Core
 {
@@ -72,6 +73,11 @@ public:
   // the JIT'ed code.
   void ClearSafe();
 
+  // DolphinQt's JITWidget needs EraseSingleBlock. Nothing else (from outside of the Core) should
+  // use it, or else JitBlockTableModel will contain a dangling reference. If something else from
+  // outside of the Core *must* use this, consider reworking the logic in JITWidget.
+  void EraseSingleBlock(const JitBlock& block);
+
   // If "forced" is true, a recompile is being requested on code that hasn't been modified.
   void InvalidateICache(u32 address, u32 size, bool forced);
   void InvalidateICacheLine(u32 address);

Source/UnitTests/Core/PageFaultTest.cpp

@@ -40,6 +40,7 @@ public:
   // JitBase methods
   JitBaseBlockCache* GetBlockCache() override { return nullptr; }
   void Jit(u32 em_address) override {}
+  void EraseSingleBlock(const JitBlock&) override {}
   const CommonAsmRoutinesBase* GetAsmRoutines() override { return nullptr; }
   virtual bool HandleFault(uintptr_t access_address, SContext* ctx) override
   {