[BOLT] Simplify RAState helpers (NFCI)

[bgergely0]

• unify isRAStateSigned and isRAStateUnsigned to a common getRAState,
• unify setRASigned and setRAUnsigned into setRAState(MCInst, bool),
• update users of these to match the new implementations.

[bgergely0]

• [BOLT] Improve InsertNegateRAStatePass::inferUnknownStates #163381
• [BOLT] Simplify RAState helpers (NFCI) #162820 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Gergely Bálint (bgergely0)

Changes

• unify isRAStateSigned and isRAStateUnsigned to a common getRAState,
• unify setRASigned and setRAUnsigned into setRAState(MCInst, bool),
• update users of these to match the new implementations.

---

Full diff: https://github.com/llvm/llvm-project/pull/162820.diff

4 Files Affected:

• (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+6-13)
• (modified) bolt/lib/Core/MCPlusBuilder.cpp (+11-16)
• (modified) bolt/lib/Passes/InsertNegateRAStatePass.cpp (+27-11)
• (modified) bolt/lib/Passes/MarkRAStates.cpp (+2-10)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 2772de73081d1..a361228354421 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1361,20 +1361,13 @@ class MCPlusBuilder {
   /// Return true if \p Inst has RestoreState annotation.
   bool hasRestoreState(const MCInst &Inst) const;
 
-  /// Stores RA Signed annotation on \p Inst.
-  void setRASigned(MCInst &Inst) const;
+  /// Sets kRASigned or kRAUnsigned annotation on \p Inst.
+  /// Fails if \p Inst has either annotation already set.
+  void setRAState(MCInst &Inst, bool State) const;
 
-  /// Return true if \p Inst has Signed RA annotation.
-  bool isRASigned(const MCInst &Inst) const;
-
-  /// Stores RA Unsigned annotation on \p Inst.
-  void setRAUnsigned(MCInst &Inst) const;
-
-  /// Return true if \p Inst has Unsigned RA annotation.
-  bool isRAUnsigned(const MCInst &Inst) const;
-
-  /// Return true if \p Inst doesn't have any annotation related to RA state.
-  bool isRAStateUnknown(const MCInst &Inst) const;
+  /// Return true if \p Inst has kRASigned annotation, false if it has
+  /// kRAUnsigned annotation, and std::nullopt if neither annotation is set.
+  std::optional<bool> getRAState(const MCInst &Inst) const;
 
   /// Return true if the instruction is a call with an exception handling info.
   virtual bool isInvoke(const MCInst &Inst) const {
diff --git a/bolt/lib/Core/MCPlusBuilder.cpp b/bolt/lib/Core/MCPlusBuilder.cpp
index e96de80bfa701..0cb4ba1ebfbd7 100644
--- a/bolt/lib/Core/MCPlusBuilder.cpp
+++ b/bolt/lib/Core/MCPlusBuilder.cpp
@@ -186,26 +186,21 @@ bool MCPlusBuilder::hasRestoreState(const MCInst &Inst) const {
   return hasAnnotation(Inst, MCAnnotation::kRestoreState);
 }
 
-void MCPlusBuilder::setRASigned(MCInst &Inst) const {
+void MCPlusBuilder::setRAState(MCInst &Inst, bool State) const {
   assert(!hasAnnotation(Inst, MCAnnotation::kRASigned));
-  setAnnotationOpValue(Inst, MCAnnotation::kRASigned, true);
-}
-
-bool MCPlusBuilder::isRASigned(const MCInst &Inst) const {
-  return hasAnnotation(Inst, MCAnnotation::kRASigned);
-}
-
-void MCPlusBuilder::setRAUnsigned(MCInst &Inst) const {
   assert(!hasAnnotation(Inst, MCAnnotation::kRAUnsigned));
-  setAnnotationOpValue(Inst, MCAnnotation::kRAUnsigned, true);
+  if (State)
+    setAnnotationOpValue(Inst, MCAnnotation::kRASigned, true);
+  else
+    setAnnotationOpValue(Inst, MCAnnotation::kRAUnsigned, true);
 }
 
-bool MCPlusBuilder::isRAUnsigned(const MCInst &Inst) const {
-  return hasAnnotation(Inst, MCAnnotation::kRAUnsigned);
-}
-
-bool MCPlusBuilder::isRAStateUnknown(const MCInst &Inst) const {
-  return !(isRAUnsigned(Inst) || isRASigned(Inst));
+std::optional<bool> MCPlusBuilder::getRAState(const MCInst &Inst) const {
+  if (hasAnnotation(Inst, MCAnnotation::kRASigned))
+    return true;
+  if (hasAnnotation(Inst, MCAnnotation::kRAUnsigned))
+    return false;
+  return std::nullopt;
 }
 
 std::optional<MCLandingPad> MCPlusBuilder::getEHInfo(const MCInst &Inst) const {
diff --git a/bolt/lib/Passes/InsertNegateRAStatePass.cpp b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
index 33664e1160a7b..c29f26c10f253 100644
--- a/bolt/lib/Passes/InsertNegateRAStatePass.cpp
+++ b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
@@ -40,6 +40,7 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
     coverFunctionFragmentStart(BF, FF);
     bool FirstIter = true;
     MCInst PrevInst;
+    bool PrevRAState = false;
     // As this pass runs after function splitting, we should only check
     // consecutive instructions inside FunctionFragments.
     for (BinaryBasicBlock *BB : FF) {
@@ -47,18 +48,22 @@ void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
         MCInst &Inst = *It;
         if (BC.MIB->isCFI(Inst))
           continue;
+        auto RAState = BC.MIB->getRAState(Inst);
+        if (!RAState) {
+          BC.errs() << "BOLT-ERROR: unknown RAState after inferUnknownStates "
+                    << " in function " << BF.getPrintName() << "\n";
+        }
         if (!FirstIter) {
           // Consecutive instructions with different RAState means we need to
           // add a OpNegateRAState.
-          if ((BC.MIB->isRASigned(PrevInst) && BC.MIB->isRAUnsigned(Inst)) ||
-              (BC.MIB->isRAUnsigned(PrevInst) && BC.MIB->isRASigned(Inst))) {
+          if (*RAState != PrevRAState)
             It = BF.addCFIInstruction(
                 BB, It, MCCFIInstruction::createNegateRAState(nullptr));
-          }
         } else {
           FirstIter = false;
         }
         PrevInst = *It;
+        PrevRAState = *RAState;
       }
     }
   }
@@ -81,10 +86,15 @@ void InsertNegateRAState::coverFunctionFragmentStart(BinaryFunction &BF,
       });
   // If a function is already split in the input, the first FF can also start
   // with Signed state. This covers that scenario as well.
-  if (BC.MIB->isRASigned(*((*FirstNonEmpty)->begin()))) {
+  auto RAState = BC.MIB->getRAState(*(*FirstNonEmpty)->begin());
+  if (!RAState) {
+    BC.errs() << "BOLT-ERROR: unknown RAState after inferUnknownStates "
+              << " in function " << BF.getPrintName() << "\n";
+    return;
+  }
+  if (*RAState)
     BF.addCFIInstruction(*FirstNonEmpty, (*FirstNonEmpty)->begin(),
                          MCCFIInstruction::createNegateRAState(nullptr));
-  }
 }
 
 void InsertNegateRAState::inferUnknownStates(BinaryFunction &BF) {
@@ -96,15 +106,21 @@ void InsertNegateRAState::inferUnknownStates(BinaryFunction &BF) {
       if (BC.MIB->isCFI(Inst))
         continue;
 
-      if (!FirstIter && BC.MIB->isRAStateUnknown(Inst)) {
-        if (BC.MIB->isRASigned(PrevInst) || BC.MIB->isPSignOnLR(PrevInst)) {
-          BC.MIB->setRASigned(Inst);
-        } else if (BC.MIB->isRAUnsigned(PrevInst) ||
-                   BC.MIB->isPAuthOnLR(PrevInst)) {
-          BC.MIB->setRAUnsigned(Inst);
+      auto RAState = BC.MIB->getRAState(Inst);
+      if (!FirstIter && !RAState) {
+        if (BC.MIB->isPSignOnLR(PrevInst))
+          RAState = true;
+        else if (BC.MIB->isPAuthOnLR(PrevInst))
+          RAState = false;
+        else {
+          auto PrevRAState = BC.MIB->getRAState(PrevInst);
+          RAState = PrevRAState ? *PrevRAState : false;
         }
+        BC.MIB->setRAState(Inst, *RAState);
       } else {
         FirstIter = false;
+        if (!RAState)
+          BC.MIB->setRAState(Inst, BF.getInitialRAState());
       }
       PrevInst = Inst;
     }
diff --git a/bolt/lib/Passes/MarkRAStates.cpp b/bolt/lib/Passes/MarkRAStates.cpp
index af6a5ca7e31e5..81d5a2257888a 100644
--- a/bolt/lib/Passes/MarkRAStates.cpp
+++ b/bolt/lib/Passes/MarkRAStates.cpp
@@ -70,9 +70,6 @@ bool MarkRAStates::runOnFunction(BinaryFunction &BF) {
           BF.setIgnored();
           return false;
         }
-        // The signing instruction itself is unsigned, the next will be
-        // signed.
-        BC.MIB->setRAUnsigned(Inst);
       } else if (BC.MIB->isPAuthOnLR(Inst)) {
         if (!RAState) {
           // RA authenticating instructions should only follow signed RA state.
@@ -83,15 +80,10 @@ bool MarkRAStates::runOnFunction(BinaryFunction &BF) {
           BF.setIgnored();
           return false;
         }
-        // The authenticating instruction itself is signed, but the next will be
-        // unsigned.
-        BC.MIB->setRASigned(Inst);
-      } else if (RAState) {
-        BC.MIB->setRASigned(Inst);
-      } else {
-        BC.MIB->setRAUnsigned(Inst);
       }
 
+      BC.MIB->setRAState(Inst, RAState);
+
       // Updating RAState. All updates are valid from the next instruction.
       // Because the same instruction can have remember and restore, the order
       // here is relevant. This is the reason to loop over Annotations instead