[abi] Add support for marking registers as unavailable in FuncFrame (#477)

src/asmjit/arm/a64rapass.cpp

@@ -654,18 +654,20 @@ void ARMRAPass::onInit() noexcept {
   _scratchRegIndexes[0] = uint8_t(27);
   _scratchRegIndexes[1] = uint8_t(28);
 
+  const FuncFrame& frame = _func->frame();
+
   // The architecture specific setup makes implicitly all registers available. So
   // make unavailable all registers that are special and cannot be used in general.
-  bool hasFP = _func->frame().hasPreservedFP();
+  bool hasFP = frame.hasPreservedFP();
 
   // Apple ABI requires that the frame-pointer register is not changed by leaf functions and properly updated
   // by non-leaf functions. So, let's make this register unavailable as it's just not safe to update it.
   if (hasFP || cc()->environment().isDarwin()) {
     makeUnavailable(RegGroup::kGp, Gp::kIdFp);
   }
-
   makeUnavailable(RegGroup::kGp, Gp::kIdSp);
   makeUnavailable(RegGroup::kGp, Gp::kIdOs); // OS-specific use, usually TLS.
+  makeUnavailable(frame._unavailableRegs);
 
   _sp = sp;
   _fp = x29;

src/asmjit/core/func.h

@@ -1194,6 +1194,8 @@ public:
   RegMasks _dirtyRegs {};
   //! Registers that must be preserved (copied from CallConv).
   RegMasks _preservedRegs {};
+  //! Registers that are unavailable.
+  RegMasks _unavailableRegs {};
   //! Size to save/restore per register group.
   Support::Array<uint8_t, Globals::kNumVirtGroups> _saveRestoreRegSize {};
   //! Alignment of save/restore area per register group.
@@ -1566,6 +1568,58 @@ public:
     return _preservedRegs[group];
   }
 
+  //! Sets which registers (as a mask) are unavailable for allocation.
+  //!
+  //! \note This completely overwrites the current unavailable mask.
+  inline void setUnavailableRegs(RegGroup group, RegMask regs) noexcept {
+    ASMJIT_ASSERT(group <= RegGroup::kMaxVirt);
+    _unavailableRegs[group] = regs;
+  }
+
+  //! Adds registers (as a mask) to the unavailable set.
+  inline void addUnavailableRegs(RegGroup group, RegMask regs) noexcept {
+    ASMJIT_ASSERT(group <= RegGroup::kMaxVirt);
+    _unavailableRegs[group] |= regs;
+  }
+
+  //! Adds a single register to the unavailable set.
+  inline void addUnavailableRegs(const BaseReg& reg) noexcept {
+    ASMJIT_ASSERT(reg.id() < Globals::kMaxPhysRegs);
+    addUnavailableRegs(reg.group(), Support::bitMask(reg.id()));
+  }
+
+  //! Adds multiple registers to the unavailable set.
+  template<typename... Args>
+  inline void addUnavailableRegs(const BaseReg& reg, Args&&... args) noexcept {
+    addUnavailableRegs(reg);
+    addUnavailableRegs(std::forward<Args>(args)...);
+  }
+
+  //! Clears all unavailable registers across all register groups (i.e., makes them all available again).
+  ASMJIT_INLINE_NODEBUG void clearUnavailableRegs() noexcept {
+    for (size_t i = 0; i < ASMJIT_ARRAY_SIZE(_unavailableRegs); i++)
+      _unavailableRegs[i] = 0;
+  }
+
+  //! Clears all unavailable registers in a specific register group.
+  inline void clearUnavailableRegs(RegGroup group) noexcept {
+    ASMJIT_ASSERT(group <= RegGroup::kMaxVirt);
+    _unavailableRegs[group] = 0;
+  }
+
+  //! Returns the set of unavailable registers for the given group.
+  [[nodiscard]]
+  inline RegMask unavailableRegs(RegGroup group) const noexcept {
+    ASMJIT_ASSERT(group <= RegGroup::kMaxVirt);
+    return _unavailableRegs[group];
+  }
+
+  //! Returns all unavailable registers as a Support::Array<>.
+  [[nodiscard]]
+  ASMJIT_INLINE_NODEBUG const RegMasks& unavailableRegs() const noexcept {
+    return _unavailableRegs;
+  }
+
   //! Returns the size of a save-restore are for the required register `group`.
   [[nodiscard]]
   inline uint32_t saveRestoreRegSize(RegGroup group) const noexcept {

src/asmjit/x86/x86rapass.cpp

@@ -1299,14 +1299,17 @@ void X86RAPass::onInit() noexcept {
   _scratchRegIndexes[0] = uint8_t(Gp::kIdCx);
   _scratchRegIndexes[1] = uint8_t(baseRegCount - 1);
 
+  const FuncFrame& frame = _func->frame();
+
   // The architecture specific setup makes implicitly all registers available. So
   // make unavailable all registers that are special and cannot be used in general.
-  bool hasFP = _func->frame().hasPreservedFP();
+  bool hasFP = frame.hasPreservedFP();
 
   makeUnavailable(RegGroup::kGp, Gp::kIdSp);   // ESP|RSP used as a stack-pointer (SP).
   if (hasFP) {
     makeUnavailable(RegGroup::kGp, Gp::kIdBp); // EBP|RBP used as a frame-pointer (FP).
   }
+  makeUnavailable(frame._unavailableRegs);
 
   _sp = cc()->zsp();
   _fp = cc()->zbp();