[BOLT][BTI] Fix assertions checking getNumOperands (#174600)

bgergely0 · web-flow · commit 76c300c8c71c · 2026-01-07T10:54:00.000+01:00
Several BTI-related functions are checking that a call MCInst has one
non-annotation operand.

This patch changes these checks to use MCPlus::getNumPrimeOperands,
instead of getNumOperands.

Testing: 
added annotations to existing gtests to serve as regression
tests. These now also explicitly check getNumOperands and getNumPrimeOperands
usage on the annotated MCInsts.
diff --git a/bolt/lib/Passes/PointerAuthCFIFixup.cpp b/bolt/lib/Passes/PointerAuthCFIFixup.cpp
@@ -183,7 +183,7 @@ void PointerAuthCFIFixup::fillUnknownStateInBB(BinaryContext &BC,
 void PointerAuthCFIFixup::markUnknownBlock(BinaryContext &BC,
                                            BinaryBasicBlock &BB, bool State) {
   // If we call this when an Instruction has either kRASigned or kRAUnsigned
-  // annotation, setRASigned or setRAUnsigned would fail.
+  // annotation, setRAState would fail.
   assert(isUnknownBlock(BC, BB) &&
          "markUnknownBlock should only be called on unknown blocks");
   for (MCInst &Inst : BB) {
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2788,9 +2788,10 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     bool CallTarget = BTI == BTIKind::C || BTI == BTIKind::JC;
     bool JumpTarget = BTI == BTIKind::J || BTI == BTIKind::JC;
     unsigned HintNum = getBTIHintNum(CallTarget, JumpTarget);
-    bool IsExplicitBTI =
-        Inst.getOpcode() == AArch64::HINT && Inst.getNumOperands() == 1 &&
-        Inst.getOperand(0).isImm() && Inst.getOperand(0).getImm() == HintNum;
+    bool IsExplicitBTI = Inst.getOpcode() == AArch64::HINT &&
+                         MCPlus::getNumPrimeOperands(Inst) == 1 &&
+                         Inst.getOperand(0).isImm() &&
+                         Inst.getOperand(0).getImm() == HintNum;
 
     // Only "BTI C" can be implicit.
     bool IsImplicitBTI =
@@ -2818,7 +2819,7 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     // 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 &&
+      assert(MCPlus::getNumPrimeOperands(Call) == 1 &&
              "Indirect branch needs to have 1 operand.");
       assert(Call.getOperand(0).isReg() &&
              "Indirect branch does not have a register operand.");
@@ -2856,7 +2857,7 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     // 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 &&
+      assert(MCPlus::getNumPrimeOperands(Call) == 1 &&
              "Indirect branch needs to have 1 operand.");
       assert(Call.getOperand(0).isReg() &&
              "Indirect branch does not have a register operand.");
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -152,13 +152,19 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   MCInst BTIjc;
   BC->MIB->createBTI(BTIjc, BTIKind::JC);
   BB->addInstruction(BTIjc);
+  BC->MIB->setRAState(BTIjc, true);
+  ASSERT_NE(BTIjc.getNumOperands(), 1u);
+  ASSERT_EQ(MCPlus::getNumPrimeOperands(BTIjc), 1u);
   auto II = BB->begin();
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 38);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, BTIKind::JC));
 
   MCInst BTIj;
   BC->MIB->createBTI(BTIj, BTIKind::J);
+  BC->MIB->setRAState(BTIj, false);
+  ASSERT_NE(BTIj.getNumOperands(), 1u);
+  ASSERT_EQ(MCPlus::getNumPrimeOperands(BTIj), 1u);
   II = BB->addInstruction(BTIj);
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 36);
@@ -206,6 +212,11 @@ TEST_P(MCPlusBuilderTester, AArch64_insertBTI_0) {
   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);
+  // Adding an annotation to the call, to check if param numbers are calculated
+  // correctly. Could be any other annotation as well.
+  BC->MIB->setRAState(CallInst, false);
+  ASSERT_NE(CallInst.getNumOperands(), 1u);
+  ASSERT_EQ(MCPlus::getNumPrimeOperands(CallInst), 1u);
   auto II = BB->begin();
   ASSERT_FALSE(BC->MIB->isCallCoveredByBTI(CallInst, *II));
   BC->MIB->insertBTI(*BB, CallInst);
@@ -223,6 +234,9 @@ TEST_P(MCPlusBuilderTester, AArch64_insertBTI_1) {
   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->setRAState(CallInst, true);
+  ASSERT_NE(CallInst.getNumOperands(), 1u);
+  ASSERT_EQ(MCPlus::getNumPrimeOperands(CallInst), 1u);
   auto II = BB->begin();
   ASSERT_TRUE(BC->MIB->isCallCoveredByBTI(CallInst, *II));
   BC->MIB->insertBTI(*BB, CallInst);
