[BOLT][BTI] Add MCPlusBuilder::addBTItoBBStartusers/bgergely0/bolt-add-bti-to-bb

This function contains most of the logic for BTI:
- it takes the BasicBlock and the instruction used to jump to it.
- then it checks if the first non-pseudo instruction is a sufficient
landing pad for the used call.
- if not, it generates the correct BTI instruction.

Also introduce the isBTIVariantCoveringCall helper to simplify the logic.

3 files changed, 193 insertions, 0 deletions

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index b4dc7b3a8b93..4f5a3f191fc2 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1888,6 +1888,19 @@ public:
     llvm_unreachable("not implemented");
   }
 
+  /// Checks if the indirect call / jump is accepted by the landing pad at the
+  /// start of the target BasicBlock.
+  virtual bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const {
+    llvm_unreachable("not implemented");
+    return false;
+  }
+
+  /// Adds a BTI landing pad to the start of the BB, that matches the indirect
+  /// call/jump inst.
+  virtual void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const {
+    llvm_unreachable("not implemented");
+  }
+
   /// Store \p Target absolute address to \p RegName
   virtual InstructionListType materializeAddress(const MCSymbol *Target,
                                                  MCContext *Ctx,
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 3f1eb7eab115..17a9634d8e6b 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2780,6 +2780,81 @@ public:
     Inst.addOperand(MCOperand::createImm(HintNum));
   }
 
+  bool isBTIVariantCoveringCall(MCInst &Call, MCInst &Pad) const override {
+    assert((isIndirectCall(Call) || isIndirectBranch(Call)) &&
+           "Not an indirect call or branch.");
+
+    // A BLR can be accepted by a BTI c.
+    if (isIndirectCall(Call))
+      return isBTILandingPad(Pad, true, false) ||
+             isBTILandingPad(Pad, true, true);
+
+    // A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
+    // x16 or x17. If the operand is not x16 or x17, it can be accepted by a BTI
+    // j or BTI jc (and not BTI c).
+    if (isIndirectBranch(Call)) {
+      assert(Call.getNumOperands() == 1 &&
+             "Indirect branch needs to have 1 operand.");
+      assert(Call.getOperand(0).isReg() &&
+             "Indirect branch does not have a register operand.");
+      MCPhysReg Reg = Call.getOperand(0).getReg();
+      if (Reg == AArch64::X16 || Reg == AArch64::X17)
+        return isBTILandingPad(Pad, true, false) ||
+               isBTILandingPad(Pad, false, true) ||
+               isBTILandingPad(Pad, true, true);
+      return isBTILandingPad(Pad, false, true) ||
+             isBTILandingPad(Pad, true, true);
+    }
+    return false;
+  }
+
+  void addBTItoBBStart(BinaryBasicBlock &BB, MCInst &Call) const override {
+    auto II = BB.getFirstNonPseudo();
+    if (II != BB.end()) {
+      if (isBTIVariantCoveringCall(Call, *II))
+        return;
+      // A BLR can be accepted by a BTI c.
+      if (isIndirectCall(Call)) {
+        // if we have a BTI j at the start, extend it to a BTI jc,
+        // otherwise insert a new BTI c.
+        if (isBTILandingPad(*II, false, true)) {
+          updateBTIVariant(*II, true, true);
+        } else {
+          MCInst BTIInst;
+          createBTI(BTIInst, true, false);
+          BB.insertInstruction(II, BTIInst);
+        }
+      }
+
+      // A BR can be accepted by a BTI j or BTI c (and BTI jc) IF the operand is
+      // x16 or x17. If the operand is not x16 or x17, it can be accepted by a
+      // BTI j or BTI jc (and not BTI c).
+      if (isIndirectBranch(Call)) {
+        assert(Call.getNumOperands() == 1 &&
+               "Indirect branch needs to have 1 operand.");
+        assert(Call.getOperand(0).isReg() &&
+               "Indirect branch does not have a register operand.");
+        MCPhysReg Reg = Call.getOperand(0).getReg();
+        if (Reg == AArch64::X16 || Reg == AArch64::X17) {
+          // Add a new BTI c
+          MCInst BTIInst;
+          createBTI(BTIInst, true, false);
+          BB.insertInstruction(II, BTIInst);
+        } else {
+          // If BB starts with a BTI c, extend it to BTI jc,
+          // otherwise insert a new BTI j.
+          if (isBTILandingPad(*II, true, false)) {
+            updateBTIVariant(*II, true, true);
+          } else {
+            MCInst BTIInst;
+            createBTI(BTIInst, false, true);
+            BB.insertInstruction(II, BTIInst);
+          }
+        }
+      }
+    }
+  }
+
   InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
                                          MCPhysReg RegName,
                                          int64_t Addend = 0) const override {
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 02ecb87b4a5e..e08ae09e7602 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -196,6 +196,111 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   ASSERT_TRUE(BC->MIB->isImplicitBTIC(*II));
 }
 
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_0) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+  BB->addInstruction(Inst);
+  // BR x16 needs BTI c or BTI j. We prefer adding a BTI c.
+  MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_1) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst BTIc;
+  BC->MIB->createBTI(BTIc, true, false);
+  BB->addInstruction(BTIc);
+  // BR x16 needs BTI c or BTI j. We have a BTI c, no change is needed.
+  MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X16);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_2) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst BTIc;
+  BC->MIB->createBTI(BTIc, true, false);
+  BB->addInstruction(BTIc);
+  // BR x5 needs BTI j
+  // we have BTI c -> extend it to BTI jc.
+  MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_3) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+  BB->addInstruction(Inst);
+  // BR x5 needs BTI j
+  MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X5);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_4) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst Inst = MCInstBuilder(AArch64::RET).addReg(AArch64::LR);
+  BB->addInstruction(Inst);
+  // BLR needs BTI c, regardless of the register used.
+  MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_5) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst BTIj;
+  BC->MIB->createBTI(BTIj, false, true);
+  BB->addInstruction(BTIj);
+  // BLR needs BTI c, regardless of the register used.
+  // We have a BTI j -> extend it to BTI jc.
+  MCInst CallInst = MCInstBuilder(AArch64::BLR).addReg(AArch64::X5);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_addBTItoBBStart_6) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+  MCInst Paciasp =
+      MCInstBuilder(AArch64::PACIASP).addReg(AArch64::LR).addReg(AArch64::SP);
+  BB->addInstruction(Paciasp);
+  // PACI(AB)SP are implicit BTI c, no change needed.
+  MCInst CallInst = MCInstBuilder(AArch64::BR).addReg(AArch64::X17);
+  BC->MIB->addBTItoBBStart(*BB, CallInst);
+  auto II = BB->begin();
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+  ASSERT_TRUE(BC->MIB->isPSignOnLR(*II));
+}
+
 TEST_P(MCPlusBuilderTester, AArch64_CmpJNE) {
   if (GetParam() != Triple::aarch64)
     GTEST_SKIP();