================
@@ -91,44 +105,214 @@ 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.
-  auto RAState = BC.MIB->getRAState(*(*FirstNonEmpty)->begin());
+  auto II = (*FirstNonEmpty)->getFirstNonPseudo();
+  auto RAState = BC.MIB->getRAState(*II);
   if (!RAState) {
     BC.errs() << "BOLT-ERROR: unknown RAState after inferUnknownStates "
               << " in function " << BF.getPrintName() << "\n";
     PassFailed = true;
     return;
   }
   if (*RAState)
-    BF.addCFIInstruction(*FirstNonEmpty, (*FirstNonEmpty)->begin(),
+    BF.addCFIInstruction(*FirstNonEmpty, II,
                          MCCFIInstruction::createNegateRAState(nullptr));
 }
 
-void InsertNegateRAState::inferUnknownStates(BinaryFunction &BF) {
+std::optional<bool>
+InsertNegateRAState::getFirstKnownRAState(BinaryContext &BC,
+                                          BinaryBasicBlock &BB) {
+  for (const MCInst &Inst : BB) {
+    if (BC.MIB->isCFI(Inst))
+      continue;
+    auto RAStateOpt = BC.MIB->getRAState(Inst);
+    if (RAStateOpt)
+      return RAStateOpt;
+  }
+  return std::nullopt;
+}
+
+void InsertNegateRAState::fillUnknownStateInBB(BinaryContext &BC,
+                                               BinaryBasicBlock &BB) {
+
+  auto First = BB.getFirstNonPseudo();
+  if (First == BB.end())
+    return;
+  // If the first instruction has unknown RAState, we should copy the first
+  // known RAState.
+  auto RAStateOpt = BC.MIB->getRAState(*First);
+  if (!RAStateOpt) {
+    auto FirstRAState = getFirstKnownRAState(BC, BB);
+    if (!FirstRAState)
+      // We fill unknown BBs later.
+      return;
+
+    BC.MIB->setRAState(*First, *FirstRAState);
+  }
+
+  // At this point we know the RAState of the first instruction,
+  // so we can propagate the RAStates to all subsequent unknown instructions.
+  MCInst Prev = *First;
+  for (auto It = BB.begin() + 1; It != BB.end(); ++It) {
+    MCInst &Inst = *It;
+    if (BC.MIB->isCFI(Inst))
+      continue;
+
+    auto PrevRAState = BC.MIB->getRAState(Prev);
+    if (!PrevRAState)
+      llvm_unreachable("Previous Instruction has no RAState.");
+
+    auto RAState = BC.MIB->getRAState(Inst);
+    if (!RAState) {
+      if (BC.MIB->isPSignOnLR(Prev))
+        PrevRAState = true;
+      else if (BC.MIB->isPAuthOnLR(Prev))
+        PrevRAState = false;
+      BC.MIB->setRAState(Inst, *PrevRAState);
+    }
+    Prev = Inst;
+  }
+}
+
+bool InsertNegateRAState::isUnknownBlock(BinaryContext &BC,
+                                         BinaryBasicBlock &BB) {
+  for (const MCInst &Inst : BB) {
+    if (BC.MIB->isCFI(Inst))
+      continue;
+    auto RAState = BC.MIB->getRAState(Inst);
+    if (RAState)
+      return false;
+  }
+  return true;
+}
+
+void InsertNegateRAState::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.
+  assert(isUnknownBlock(BC, BB) &&
+         "markUnknownBlock should only be called on unknown blocks");
+  for (MCInst &Inst : BB) {
+    if (BC.MIB->isCFI(Inst))
+      continue;
+    BC.MIB->setRAState(Inst, State);
+  }
+}
+
+std::optional<bool> InsertNegateRAState::getRAStateByCFG(BinaryBasicBlock &BB,
+                                                         BinaryFunction &BF) {
   BinaryContext &BC = BF.getBinaryContext();
-  bool FirstIter = true;
-  MCInst PrevInst;
-  for (BinaryBasicBlock &BB : BF) {
-    for (MCInst &Inst : BB) {
-      if (BC.MIB->isCFI(Inst))
+
+  auto checkRAState = [&](std::optional<bool> &NeighborRAState, MCInst &Inst) {
+    auto RAState = BC.MIB->getRAState(Inst);
+    if (!RAState)
+      return;
+    if (!NeighborRAState) {
+      NeighborRAState = *RAState;
+      return;
+    }
+    if (NeighborRAState != *RAState) {
+      BC.outs() << "BOLT-WARNING: Conflicting RAState found in function "
+                << BF.getPrintName() << ". Function will not be optimized.\n";
+      BF.setIgnored();
+    }
+  };
+
+  // Holds the first found RAState from CFG neighbors.
+  std::optional<bool> NeighborRAState = std::nullopt;
+  if (BB.pred_size() != 0) {
+    for (BinaryBasicBlock *PredBB : BB.predecessors()) {
+      //  find last inst of Predecessor with known RA State.
+      auto LI = PredBB->getLastNonPseudo();
+      if (LI == PredBB->rend())
+        continue;
+      MCInst &LastInst = *LI;
+      checkRAState(NeighborRAState, LastInst);
+    }
+  } else if (BB.succ_size() != 0) {
+    for (BinaryBasicBlock *SuccBB : BB.successors()) {
+      //  find first inst of Successor with known RA State.
+      auto FI = SuccBB->getFirstNonPseudo();
+      if (FI == SuccBB->end())
         continue;
+      MCInst &FirstInst = *FI;
+      checkRAState(NeighborRAState, FirstInst);
+    }
+  } else {
+    llvm_unreachable("Called getRAStateByCFG on a BB with no preds or succs.");
----------------
bgergely0 wrote:

TODO: change to fatal BOLT error, based on the discussion in #162820

https://github.com/llvm/llvm-project/pull/163381
_______________________________________________
llvm-branch-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits

Reply via email to