[BOLT][BTI] Add MCPlusBuilder::updateBTIVariant

Checks if an instruction is BTI, and updates the immediate value to the
newly requested variant.

Author: bgergely0

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -1882,6 +1882,12 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
+  /// Update operand of BTI instruction.
+  virtual void updateBTIVariant(MCInst &Inst, bool CallTarget,
+                                bool JumpTarget) const {
+    llvm_unreachable("not implemented");
+  }
+
   /// Store \p Target absolute address to \p RegName
   virtual InstructionListType materializeAddress(const MCSymbol *Target,
                                                  MCContext *Ctx,

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -2774,6 +2774,14 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   }
 
 
+  void updateBTIVariant(MCInst &Inst, bool CallTarget,
+                        bool JumpTarget) const override {
+    assert(Inst.getOpcode() == AArch64::HINT && "Not a BTI instruction.");
+    unsigned HintNum = getBTIHintNum(CallTarget, JumpTarget);
+    Inst.clear();
+    Inst.addOperand(MCOperand::createImm(HintNum));
+  }
+
   InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
                                          MCPhysReg RegName,
                                          int64_t Addend = 0) const override {

bolt/unittests/Core/MCPlusBuilder.cpp

@@ -156,20 +156,26 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 38);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+  BC->MIB->updateBTIVariant(*II, true, false);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
 
   MCInst BTIj;
   BC->MIB->createBTI(BTIj, false, true);
   II = BB->addInstruction(BTIj);
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 36);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
+  BC->MIB->updateBTIVariant(*II, true, true);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
 
   MCInst BTIc;
   BC->MIB->createBTI(BTIc, true, false);
   II = BB->addInstruction(BTIc);
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 34);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+  BC->MIB->updateBTIVariant(*II, false, true);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
 
   MCInst BTIinvalid;
   ASSERT_DEATH(BC->MIB->createBTI(BTIinvalid, false, false),