[BOLT] Review changesusers/bgergely0/bolt-infer-unknown-states

2 files changed, 46 insertions, 39 deletions

diff --git a/bolt/include/bolt/Passes/InsertNegateRAStatePass.h b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h
index b4b428207b65..c819b1c4914b 100644
--- a/bolt/include/bolt/Passes/InsertNegateRAStatePass.h
+++ b/bolt/include/bolt/Passes/InsertNegateRAStatePass.h
@@ -30,6 +30,8 @@ public:
 private:
   /// Because states are tracked as MCAnnotations on individual instructions,
   /// newly inserted instructions do not have a state associated with them.
+  /// Uses fillUnknownStateInBB, fillUnknownBlocksInCFG and fillUnknownStubs in
+  /// this order of priority.
   void inferUnknownStates(BinaryFunction &BF);
 
   /// Simple case: copy RAStates to unknown insts from previous inst.
diff --git a/bolt/lib/Passes/InsertNegateRAStatePass.cpp b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
index 1b1fbf7f20e9..be9bb3bc6e22 100644
--- a/bolt/lib/Passes/InsertNegateRAStatePass.cpp
+++ b/bolt/lib/Passes/InsertNegateRAStatePass.cpp
@@ -157,9 +157,10 @@ void InsertNegateRAState::fillUnknownStateInBB(BinaryContext &BC,
     if (BC.MIB->isCFI(Inst))
       continue;
 
+    // No need to check for nullopt: we only entered this loop after the first
+    // instruction had its RAState set, and RAState is always set for the
+    // previous instruction in the previous iteration of the loop.
     auto PrevRAState = BC.MIB->getRAState(Prev);
-    if (!PrevRAState)
-      llvm_unreachable("Previous Instruction has no RAState.");
 
     auto RAState = BC.MIB->getRAState(Inst);
     if (!RAState) {
@@ -198,6 +199,47 @@ void InsertNegateRAState::markUnknownBlock(BinaryContext &BC,
   }
 }
 
+void InsertNegateRAState::fillUnknownStubs(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+  bool FirstIter = true;
+  MCInst PrevInst;
+  for (FunctionFragment &FF : BF.getLayout().fragments()) {
+    for (BinaryBasicBlock *BB : FF) {
+      if (!FirstIter && isUnknownBlock(BC, *BB)) {
+        // If we have no predecessors or successors, current BB is a Stub called
+        // from another BinaryFunction. As of #160989, we have to copy the
+        // PrevInst's RAState, because CFIs are already incorrect here.
+        if (BB->pred_size() == 0 && BB->succ_size() == 0) {
+          auto PrevRAState = BC.MIB->getRAState(PrevInst);
+          if (!PrevRAState) {
+            llvm_unreachable(
+                "Previous Instruction has no RAState in fillUnknownStubs.");
+            continue;
+          }
+
+          if (BC.MIB->isPSignOnLR(PrevInst)) {
+            PrevRAState = true;
+          } else if (BC.MIB->isPAuthOnLR(PrevInst)) {
+            PrevRAState = false;
+          }
+          markUnknownBlock(BC, *BB, *PrevRAState);
+        }
+      }
+      if (FirstIter) {
+        FirstIter = false;
+        if (isUnknownBlock(BC, *BB))
+          // Set block to unsigned, because this is the first block of the
+          // function, and all instruction in it are new instructions generated
+          // by BOLT optimizations.
+          markUnknownBlock(BC, *BB, false);
+      }
+      auto Last = BB->getLastNonPseudo();
+      if (Last != BB->rend())
+        PrevInst = *Last;
+    }
+  }
+}
+
 std::optional<bool> InsertNegateRAState::getRAStateByCFG(BinaryBasicBlock &BB,
                                                          BinaryFunction &BF) {
   BinaryContext &BC = BF.getBinaryContext();
@@ -244,43 +286,6 @@ std::optional<bool> InsertNegateRAState::getRAStateByCFG(BinaryBasicBlock &BB,
   return NeighborRAState;
 }
 
-void InsertNegateRAState::fillUnknownStubs(BinaryFunction &BF) {
-  BinaryContext &BC = BF.getBinaryContext();
-  bool FirstIter = true;
-  MCInst PrevInst;
-  for (FunctionFragment &FF : BF.getLayout().fragments()) {
-    for (BinaryBasicBlock *BB : FF) {
-      if (!FirstIter && isUnknownBlock(BC, *BB)) {
-        // If we have no predecessors or successors, current BB is a Stub called
-        // from another BinaryFunction. As of #160989, we have to copy the
-        // PrevInst's RAState, because CFIs are already incorrect here.
-        if (BB->pred_size() == 0 && BB->succ_size() == 0) {
-          auto PrevRAState = BC.MIB->getRAState(PrevInst);
-          if (!PrevRAState) {
-            llvm_unreachable(
-                "Previous Instruction has no RAState in fillUnknownStubs.");
-            continue;
-          }
-
-          if (BC.MIB->isPSignOnLR(PrevInst)) {
-            PrevRAState = true;
-          } else if (BC.MIB->isPAuthOnLR(PrevInst)) {
-            PrevRAState = false;
-          }
-          markUnknownBlock(BC, *BB, *PrevRAState);
-        }
-      }
-      if (FirstIter) {
-        FirstIter = false;
-        if (isUnknownBlock(BC, *BB))
-          markUnknownBlock(BC, *BB, false);
-      }
-      auto Last = BB->getLastNonPseudo();
-      if (Last != BB->rend())
-        PrevInst = *Last;
-    }
-  }
-}
 void InsertNegateRAState::fillUnknownBlocksInCFG(BinaryFunction &BF) {
   BinaryContext &BC = BF.getBinaryContext();