JitRegCache: RCForkGuard

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

@@ -970,7 +970,7 @@ BitSet8 Jit64::ComputeStaticGQRs(const PPCAnalyst::CodeBlock& cb) const
 BitSet32 Jit64::CallerSavedRegistersInUse() const
 {
   BitSet32 result;
-  for (size_t i = 0; i < RegCache::NUM_XREGS; i++)
+  for (size_t i = 0; i < NUM_XREGS; i++)
   {
     if (!gpr.IsFreeX(i))
       result[i] = true;

Source/Core/Core/PowerPC/Jit64/RegCache/CachedReg.h

@@ -132,6 +132,7 @@ private:
 class RCConstraint
 {
 public:
+  bool IsActive() const { return realized || bind || write || read || kill_imm; }
   bool IsRealized() const { return realized; }
 
   bool ShouldBind() const { return bind; }

Source/Core/Core/PowerPC/Jit64/RegCache/JitRegCache.cpp

@@ -267,6 +267,28 @@ void RCX64Reg::Unlock()
   contents = std::monostate{};
 }
 
+RCForkGuard::RCForkGuard(RegCache& rc_) : rc(&rc_), m_regs(rc_.m_regs), m_xregs(rc_.m_xregs)
+{
+  ASSERT(!rc->IsAnyConstraintActive());
+}
+
+RCForkGuard::RCForkGuard(RCForkGuard&& other) noexcept
+    : rc(other.rc), m_regs(std::move(other.m_regs)), m_xregs(std::move(other.m_xregs))
+{
+  other.rc = nullptr;
+}
+
+void RCForkGuard::EndFork()
+{
+  if (!rc)
+    return;
+
+  ASSERT(!rc->IsAnyConstraintActive());
+  rc->m_regs = m_regs;
+  rc->m_xregs = m_xregs;
+  rc = nullptr;
+}
+
 RegCache::RegCache(Jit64& jit) : m_jit{jit}
 {
 }
@@ -595,6 +617,11 @@ RCX64Reg RegCache::Scratch(X64Reg xr)
   return RCX64Reg{this, xr};
 }
 
+RCForkGuard RegCache::Fork()
+{
+  return RCForkGuard{*this};
+}
+
 void RegCache::NewLock(preg_t preg)
 {
   m_regs[preg].Lock();
@@ -663,3 +690,9 @@ void RegCache::Realize(preg_t preg)
 
   m_constraints[preg].Realized();
 }
+
+bool RegCache::IsAnyConstraintActive() const
+{
+  return std::any_of(m_constraints.begin(), m_constraints.end(),
+                     [](const auto& c) { return c.IsActive(); });
+}

Source/Core/Core/PowerPC/Jit64/RegCache/JitRegCache.h

@@ -22,6 +22,7 @@ class RCX64Reg;
 class RegCache;
 
 using preg_t = size_t;
+static constexpr size_t NUM_XREGS = 16;
 
 class RCOpArg
 {
@@ -97,6 +98,28 @@ private:
   std::variant<std::monostate, Gen::X64Reg, preg_t> contents;
 };
 
+class RCForkGuard
+{
+public:
+  ~RCForkGuard() { EndFork(); }
+  RCForkGuard(RCForkGuard&&) noexcept;
+
+  RCForkGuard(const RCForkGuard&) = delete;
+  RCForkGuard& operator=(const RCForkGuard&) = delete;
+  RCForkGuard& operator=(RCForkGuard&&) = delete;
+
+  void EndFork();
+
+private:
+  friend class RegCache;
+
+  explicit RCForkGuard(RegCache& rc_);
+
+  RegCache* rc;
+  std::array<PPCCachedReg, 32> m_regs;
+  std::array<X64CachedReg, NUM_XREGS> m_xregs;
+};
+
 class RegCache
 {
 public:
@@ -106,8 +129,6 @@ public:
     MaintainState,
   };
 
-  static constexpr size_t NUM_XREGS = 16;
-
   explicit RegCache(Jit64& jit);
   virtual ~RegCache() = default;
 
@@ -216,9 +237,12 @@ public:
   RCX64Reg Bind(preg_t preg, RCMode mode);
   RCX64Reg Scratch(Gen::X64Reg xr);
 
+  RCForkGuard Fork();
+
 protected:
   friend class RCOpArg;
   friend class RCX64Reg;
+  friend class RCForkGuard;
 
   virtual void StoreRegister(preg_t preg, const Gen::OpArg& new_loc) = 0;
   virtual void LoadRegister(preg_t preg, Gen::X64Reg new_loc) = 0;
@@ -240,6 +264,8 @@ protected:
   bool IsRealized(preg_t preg) const;
   void Realize(preg_t preg);
 
+  bool IsAnyConstraintActive() const;
+
   Jit64& m_jit;
   std::array<PPCCachedReg, 32> m_regs;
   std::array<X64CachedReg, NUM_XREGS> m_xregs;