Revert "[BOLT] Refactor MCInstReference and move it to Core (NFC)"

bolt/include/bolt/Passes/PAuthGadgetScanner.h

@@ -11,13 +11,187 @@
 
 #include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/BinaryFunction.h"
-#include "bolt/Core/MCInstUtils.h"
 #include "bolt/Passes/BinaryPasses.h"
 #include "llvm/Support/raw_ostream.h"
 #include <memory>
 
 namespace llvm {
 namespace bolt {
+
+/// @brief  MCInstReference represents a reference to an MCInst as stored either
+/// in a BinaryFunction (i.e. before a CFG is created), or in a BinaryBasicBlock
+/// (after a CFG is created). It aims to store the necessary information to be
+/// able to find the specific MCInst in either the BinaryFunction or
+/// BinaryBasicBlock data structures later, so that e.g. the InputAddress of
+/// the corresponding instruction can be computed.
+
+struct MCInstInBBReference {
+  BinaryBasicBlock *BB;
+  int64_t BBIndex;
+  MCInstInBBReference(BinaryBasicBlock *BB, int64_t BBIndex)
+      : BB(BB), BBIndex(BBIndex) {}
+  MCInstInBBReference() : BB(nullptr), BBIndex(0) {}
+  static MCInstInBBReference get(const MCInst *Inst, BinaryFunction &BF) {
+    for (BinaryBasicBlock &BB : BF)
+      for (size_t I = 0; I < BB.size(); ++I)
+        if (Inst == &BB.getInstructionAtIndex(I))
+          return MCInstInBBReference(&BB, I);
+    return {};
+  }
+  bool operator==(const MCInstInBBReference &RHS) const {
+    return BB == RHS.BB && BBIndex == RHS.BBIndex;
+  }
+  bool operator<(const MCInstInBBReference &RHS) const {
+    return std::tie(BB, BBIndex) < std::tie(RHS.BB, RHS.BBIndex);
+  }
+  operator MCInst &() const {
+    assert(BB != nullptr);
+    return BB->getInstructionAtIndex(BBIndex);
+  }
+  uint64_t getAddress() const {
+    // 4 bytes per instruction on AArch64.
+    // FIXME: the assumption of 4 byte per instruction needs to be fixed before
+    // this method gets used on any non-AArch64 binaries (but should be fine for
+    // pac-ret analysis, as that is an AArch64-specific feature).
+    return BB->getFunction()->getAddress() + BB->getOffset() + BBIndex * 4;
+  }
+};
+
+raw_ostream &operator<<(raw_ostream &OS, const MCInstInBBReference &);
+
+struct MCInstInBFReference {
+  BinaryFunction *BF;
+  uint64_t Offset;
+  MCInstInBFReference(BinaryFunction *BF, uint64_t Offset)
+      : BF(BF), Offset(Offset) {}
+
+  static MCInstInBFReference get(const MCInst *Inst, BinaryFunction &BF) {
+    for (auto &I : BF.instrs())
+      if (Inst == &I.second)
+        return MCInstInBFReference(&BF, I.first);
+    return {};
+  }
+
+  MCInstInBFReference() : BF(nullptr), Offset(0) {}
+  bool operator==(const MCInstInBFReference &RHS) const {
+    return BF == RHS.BF && Offset == RHS.Offset;
+  }
+  bool operator<(const MCInstInBFReference &RHS) const {
+    return std::tie(BF, Offset) < std::tie(RHS.BF, RHS.Offset);
+  }
+  operator MCInst &() const {
+    assert(BF != nullptr);
+    return *BF->getInstructionAtOffset(Offset);
+  }
+
+  uint64_t getOffset() const { return Offset; }
+
+  uint64_t getAddress() const { return BF->getAddress() + getOffset(); }
+};
+
+raw_ostream &operator<<(raw_ostream &OS, const MCInstInBFReference &);
+
+struct MCInstReference {
+  enum Kind { FunctionParent, BasicBlockParent };
+  Kind ParentKind;
+  union U {
+    MCInstInBBReference BBRef;
+    MCInstInBFReference BFRef;
+    U(MCInstInBBReference BBRef) : BBRef(BBRef) {}
+    U(MCInstInBFReference BFRef) : BFRef(BFRef) {}
+  } U;
+  MCInstReference(MCInstInBBReference BBRef)
+      : ParentKind(BasicBlockParent), U(BBRef) {}
+  MCInstReference(MCInstInBFReference BFRef)
+      : ParentKind(FunctionParent), U(BFRef) {}
+  MCInstReference(BinaryBasicBlock *BB, int64_t BBIndex)
+      : MCInstReference(MCInstInBBReference(BB, BBIndex)) {}
+  MCInstReference(BinaryFunction *BF, uint32_t Offset)
+      : MCInstReference(MCInstInBFReference(BF, Offset)) {}
+
+  static MCInstReference get(const MCInst *Inst, BinaryFunction &BF) {
+    if (BF.hasCFG())
+      return MCInstInBBReference::get(Inst, BF);
+    return MCInstInBFReference::get(Inst, BF);
+  }
+
+  bool operator<(const MCInstReference &RHS) const {
+    if (ParentKind != RHS.ParentKind)
+      return ParentKind < RHS.ParentKind;
+    switch (ParentKind) {
+    case BasicBlockParent:
+      return U.BBRef < RHS.U.BBRef;
+    case FunctionParent:
+      return U.BFRef < RHS.U.BFRef;
+    }
+    llvm_unreachable("");
+  }
+
+  bool operator==(const MCInstReference &RHS) const {
+    if (ParentKind != RHS.ParentKind)
+      return false;
+    switch (ParentKind) {
+    case BasicBlockParent:
+      return U.BBRef == RHS.U.BBRef;
+    case FunctionParent:
+      return U.BFRef == RHS.U.BFRef;
+    }
+    llvm_unreachable("");
+  }
+
+  operator MCInst &() const {
+    switch (ParentKind) {
+    case BasicBlockParent:
+      return U.BBRef;
+    case FunctionParent:
+      return U.BFRef;
+    }
+    llvm_unreachable("");
+  }
+
+  operator bool() const {
+    switch (ParentKind) {
+    case BasicBlockParent:
+      return U.BBRef.BB != nullptr;
+    case FunctionParent:
+      return U.BFRef.BF != nullptr;
+    }
+    llvm_unreachable("");
+  }
+
+  uint64_t getAddress() const {
+    switch (ParentKind) {
+    case BasicBlockParent:
+      return U.BBRef.getAddress();
+    case FunctionParent:
+      return U.BFRef.getAddress();
+    }
+    llvm_unreachable("");
+  }
+
+  BinaryFunction *getFunction() const {
+    switch (ParentKind) {
+    case FunctionParent:
+      return U.BFRef.BF;
+    case BasicBlockParent:
+      return U.BBRef.BB->getFunction();
+    }
+    llvm_unreachable("");
+  }
+
+  BinaryBasicBlock *getBasicBlock() const {
+    switch (ParentKind) {
+    case FunctionParent:
+      return nullptr;
+    case BasicBlockParent:
+      return U.BBRef.BB;
+    }
+    llvm_unreachable("");
+  }
+};
+
+raw_ostream &operator<<(raw_ostream &OS, const MCInstReference &);
+
 namespace PAuthGadgetScanner {
 
 // The report classes are designed to be used in an immutable manner.

bolt/lib/Core/CMakeLists.txt

@@ -32,7 +32,6 @@ add_llvm_library(LLVMBOLTCore
   GDBIndex.cpp
   HashUtilities.cpp
   JumpTable.cpp
-  MCInstUtils.cpp
   MCPlusBuilder.cpp
   ParallelUtilities.cpp
   Relocation.cpp