[llvm] [BOLT][AArch64] Support for pointer authentication (PR #117578)

[Paschalis Mpeis via llvm-commits]

================
@@ -0,0 +1,241 @@
+#include "bolt/Passes/InsertNegateRAStatePass.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Utils/CommandLineOpts.h"
+#include <cstdlib>
+#include <fstream>
+#include <iterator>
+
+using namespace llvm;
+
+namespace llvm {
+namespace bolt {
+
+void InsertNegateRAState::runOnFunction(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+
+  if (BF.getState() == BinaryFunction::State::Empty) {
+    return;
+  }
+
+  if (BF.getState() != BinaryFunction::State::CFG &&
+      BF.getState() != BinaryFunction::State::CFG_Finalized) {
+    BC.outs() << "BOLT-INFO: No CFG for " << BF.getPrintName()
+              << " in InsertNegateRAStatePass\n";
+    return;
+  }
+
+  if (BF.isIgnored())
+    return;
+
+  if (!addNegateRAStateAfterPacOrAuth(BF)) {
+    // none inserted, function doesn't need more work
+    return;
+  }
+
+  auto FirstBB = BF.begin();
+  explore_call_graph(BC, &(*FirstBB));
+
+  // We have to do the walk again, starting from any undiscovered autiasp
+  // instructions, because some autiasp might not be reachable because of
+  // indirect branches but we know that autiasp block should have a Signed
+  // state, so we can work out other Unkown states starting from these nodes.
+  for (BinaryBasicBlock &BB : BF) {
+    if (BBhasAUTH(BC, &BB) && BB.isRAStateUnknown()) {
+      BB.setRASigned();
+      explore_call_graph(BC, &BB);
+    }
+  }
+
+  // insert negateRAState-s where there is a State boundary:
+  // that is: two consecutive BBs have different RA State
+  BinaryFunction::iterator PrevBB;
+  bool FirstIter = true;
+  for (auto BB = BF.begin(); BB != BF.end(); ++BB) {
+    if (!FirstIter) {
+      if ((PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           (*BB).RAState == BinaryBasicBlock::RAStateEnum::Unsigned &&
+           !BBhasAUTH(BC, &(*PrevBB))) ||
+          (PrevBB->RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           (*BB).RAState == BinaryBasicBlock::RAStateEnum::Signed &&
+           BBhasAUTH(BC, &(*PrevBB)))) {
+        auto InstRevIter = PrevBB->getLastNonPseudo();
+        MCInst LastNonPseudo = *InstRevIter;
+        auto InstIter = InstRevIter.base();
+        BF.addCFIInstruction(&(*PrevBB), InstIter,
+                             MCCFIInstruction::createNegateRAState(nullptr));
+      }
+    } else {
+      FirstIter = false;
+    }
+    PrevBB = BB;
+  }
+}
+
+void InsertNegateRAState::explore_call_graph(BinaryContext &BC,
+                                             BinaryBasicBlock *BB) {
+  std::stack<BinaryBasicBlock *> SignedStack;
+  std::stack<BinaryBasicBlock *> UnsignedStack;
+
+  // start according to the first BB
+  if (BBhasSIGN(BC, BB)) {
+    SignedStack.push(BB);
+    process_signed_BB(BC, BB, &SignedStack, &UnsignedStack);
+  } else {
+    UnsignedStack.push(BB);
+    process_unsigned_BB(BC, BB, &SignedStack, &UnsignedStack);
+  }
+
+  while (!(SignedStack.empty() && UnsignedStack.empty())) {
----------------
paschalis-mpeis wrote:

(just a suggestion; feel free to ignore)

If I'm reading this correctly, when either of this is not empty, it processes first the signed stack, followed by the unsigned one.

If so, could re-write to two loops and drop the if/else, like below?
```cpp
while (!SignedStack.empty()) { ... }
while (!UnsignedStack.empty()) { ... }
```

https://github.com/llvm/llvm-project/pull/117578