[𝘀𝗽𝗿] initial version

Created using spr 1.3.4

Author: aaupov

bolt/include/bolt/Core/BinaryContext.h

@@ -435,7 +435,18 @@ class BinaryContext {
 
   /// Return size of an entry for the given jump table \p Type.
   uint64_t getJumpTableEntrySize(JumpTable::JumpTableType Type) const {
-    return Type == JumpTable::JTT_PIC ? 4 : AsmInfo->getCodePointerSize();
+    switch (Type) {
+    case JumpTable::JTT_X86_64_PIC4:
+      return 4;
+    case JumpTable::JTT_X86_64_ABS:
+      return AsmInfo->getCodePointerSize();
+    case JumpTable::JTT_AARCH64_REL1:
+      return 1;
+    case JumpTable::JTT_AARCH64_REL2:
+      return 2;
+    case JumpTable::JTT_AARCH64_REL4:
+      return 4;
+    }
   }
 
   /// Return JumpTable containing a given \p Address.
@@ -573,14 +584,13 @@ class BinaryContext {
   /// If \p NextJTAddress is different from zero, it is used as an upper
   /// bound for jump table memory layout.
   ///
-  /// Optionally, populate \p Address from jump table entries. The entries
-  /// could be partially populated if the jump table detection fails.
+  /// If \p JT is set, populate it with jump table entries. The entries could be
+  /// partially populated if the jump table detection fails.
   bool analyzeJumpTable(const uint64_t Address,
                         const JumpTable::JumpTableType Type,
                         const BinaryFunction &BF,
                         const uint64_t NextJTAddress = 0,
-                        JumpTable::AddressesType *EntriesAsAddress = nullptr,
-                        bool *HasEntryInFragment = nullptr) const;
+                        JumpTable *JT = nullptr) const;
 
   /// After jump table locations are established, this function will populate
   /// their EntriesAsAddress based on memory contents.

bolt/include/bolt/Core/JumpTable.h

@@ -40,17 +40,42 @@ class BinaryFunction;
 /// a different label at a different offset in this jump table.
 class JumpTable : public BinaryData {
   friend class BinaryContext;
+  friend class JumpTableInfoReader;
 
   JumpTable() = delete;
   JumpTable(const JumpTable &) = delete;
   JumpTable &operator=(const JumpTable &) = delete;
 
 public:
   enum JumpTableType : char {
-    JTT_NORMAL,
-    JTT_PIC,
+    JTT_X86_64_FIRST = 0,
+    JTT_X86_64_ABS = JTT_X86_64_FIRST,
+    JTT_X86_64_PIC4,
+    JTT_X86_64_LAST = JTT_X86_64_PIC4,
+    JTT_AARCH64_FIRST,
+    JTT_AARCH64_REL1 = JTT_AARCH64_FIRST,
+    JTT_AARCH64_REL2,
+    JTT_AARCH64_REL4,
+    JTT_AARCH64_LAST = JTT_AARCH64_REL4
   };
 
+  static StringRef getTypeStr(JumpTableType Type) {
+    switch (Type) {
+    case JTT_X86_64_ABS:
+      return "X86_64_ABS";
+    case JTT_X86_64_PIC4:
+      return "X86_64_PIC4";
+    case JTT_AARCH64_REL1:
+      return "AARCH64_REL1";
+    case JTT_AARCH64_REL2:
+      return "AARCH64_REL2";
+    case JTT_AARCH64_REL4:
+      return "AARCH64_REL4";
+    }
+  }
+
+  const StringRef getTypeStr() { return getTypeStr(Type); }
+
   /// Branch statistics for jump table entries.
   struct JumpInfo {
     uint64_t Mispreds{0};
@@ -92,6 +117,16 @@ class JumpTable : public BinaryData {
   /// BinaryFunction this jump tables belongs to.
   SmallVector<BinaryFunction *, 1> Parents;
 
+  ///
+  /// AArch64-specific fields
+  ///
+
+  /// Entries are offsets relative to an arbitrary function location.
+  uint64_t BaseAddress{0};
+
+  /// Address of the instruction referencing the jump table (MemLocInstr).
+  uint64_t MemLocInstrAddress{0};
+
 private:
   /// Constructor should only be called by a BinaryContext.
   JumpTable(MCSymbol &Symbol, uint64_t Address, size_t EntrySize,

bolt/include/bolt/Rewrite/MetadataManager.h

@@ -31,6 +31,10 @@ class MetadataManager {
   /// Run initializers after sections are discovered.
   void runSectionInitializers();
 
+  /// Execute metadata initializers when functions are discovered but not yet
+  /// disassembled.
+  void runInitializersPreDisasm();
+
   /// Execute initialization of rewriters while functions are disassembled, but
   /// CFG is not yet built.
   void runInitializersPreCFG();

bolt/include/bolt/Rewrite/MetadataRewriter.h

@@ -49,6 +49,10 @@ class MetadataRewriter {
   /// but before functions are discovered.
   virtual Error sectionInitializer() { return Error::success(); }
 
+  /// Run initialization after the functions are identified but not yet
+  /// disassembled.
+  virtual Error preDisasmInitializer() { return Error::success(); }
+
   /// Interface for modifying/annotating functions in the binary based on the
   /// contents of the section. Functions are in pre-cfg state.
   virtual Error preCFGInitializer() { return Error::success(); }

bolt/include/bolt/Rewrite/MetadataRewriters.h

@@ -27,6 +27,8 @@ std::unique_ptr<MetadataRewriter> createPseudoProbeRewriter(BinaryContext &);
 
 std::unique_ptr<MetadataRewriter> createSDTRewriter(BinaryContext &);
 
+std::unique_ptr<MetadataRewriter> createJumpTableInfoReader(BinaryContext &);
+
 } // namespace bolt
 } // namespace llvm
 

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -181,6 +181,9 @@ class RewriteInstance {
   /// Process metadata in sections before functions are discovered.
   void processSectionMetadata();
 
+  /// Process metadata in special sections before functions are disassembled.
+  void processMetadataPreDisasm();
+
   /// Process metadata in special sections before CFG is built for functions.
   void processMetadataPreCFG();
 

bolt/lib/Core/BinaryContext.cpp

@@ -496,7 +496,7 @@ BinaryContext::handleAddressRef(uint64_t Address, BinaryFunction &BF,
     const MemoryContentsType MemType = analyzeMemoryAt(Address, BF);
     if (MemType == MemoryContentsType::POSSIBLE_PIC_JUMP_TABLE && IsPCRel) {
       const MCSymbol *Symbol =
-          getOrCreateJumpTable(BF, Address, JumpTable::JTT_PIC);
+          getOrCreateJumpTable(BF, Address, JumpTable::JTT_X86_64_PIC4);
 
       return std::make_pair(Symbol, 0);
     }
@@ -540,10 +540,10 @@ MemoryContentsType BinaryContext::analyzeMemoryAt(uint64_t Address,
 
   // Start with checking for PIC jump table. We expect non-PIC jump tables
   // to have high 32 bits set to 0.
-  if (analyzeJumpTable(Address, JumpTable::JTT_PIC, BF))
+  if (analyzeJumpTable(Address, JumpTable::JTT_X86_64_PIC4, BF))
     return MemoryContentsType::POSSIBLE_PIC_JUMP_TABLE;
 
-  if (analyzeJumpTable(Address, JumpTable::JTT_NORMAL, BF))
+  if (analyzeJumpTable(Address, JumpTable::JTT_X86_64_ABS, BF))
     return MemoryContentsType::POSSIBLE_JUMP_TABLE;
 
   return MemoryContentsType::UNKNOWN;
@@ -553,8 +553,7 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
                                      const JumpTable::JumpTableType Type,
                                      const BinaryFunction &BF,
                                      const uint64_t NextJTAddress,
-                                     JumpTable::AddressesType *EntriesAsAddress,
-                                     bool *HasEntryInFragment) const {
+                                     JumpTable *JT) const {
   // Target address of __builtin_unreachable.
   const uint64_t UnreachableAddress = BF.getAddress() + BF.getSize();
 
@@ -571,11 +570,11 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
   size_t TrimmedSize = 0;
 
   auto addEntryAddress = [&](uint64_t EntryAddress, bool Unreachable = false) {
-    if (!EntriesAsAddress)
+    if (!JT)
       return;
-    EntriesAsAddress->emplace_back(EntryAddress);
+    JT->EntriesAsAddress.emplace_back(EntryAddress);
     if (!Unreachable)
-      TrimmedSize = EntriesAsAddress->size();
+      TrimmedSize = JT->EntriesAsAddress.size();
   };
 
   ErrorOr<const BinarySection &> Section = getSectionForAddress(Address);
@@ -594,37 +593,48 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
   if (NextJTAddress)
     UpperBound = std::min(NextJTAddress, UpperBound);
 
-  LLVM_DEBUG({
-    using JTT = JumpTable::JumpTableType;
-    dbgs() << formatv("BOLT-DEBUG: analyzeJumpTable @{0:x} in {1}, JTT={2}\n",
-                      Address, BF.getPrintName(),
-                      Type == JTT::JTT_PIC ? "PIC" : "Normal");
-  });
+  LLVM_DEBUG(
+      dbgs() << formatv("BOLT-DEBUG: analyzeJumpTable @{0:x} in {1}, JTT={2}\n",
+                        Address, BF, JumpTable::getTypeStr(Type)));
   const uint64_t EntrySize = getJumpTableEntrySize(Type);
   for (uint64_t EntryAddress = Address; EntryAddress <= UpperBound - EntrySize;
        EntryAddress += EntrySize) {
     LLVM_DEBUG(dbgs() << "  * Checking 0x" << Twine::utohexstr(EntryAddress)
                       << " -> ");
     // Check if there's a proper relocation against the jump table entry.
     if (HasRelocations) {
-      if (Type == JumpTable::JTT_PIC &&
+      if (Type == JumpTable::JTT_X86_64_PIC4 &&
           !DataPCRelocations.count(EntryAddress)) {
         LLVM_DEBUG(
             dbgs() << "FAIL: JTT_PIC table, no relocation for this address\n");
         break;
       }
-      if (Type == JumpTable::JTT_NORMAL && !getRelocationAt(EntryAddress)) {
+      if (Type == JumpTable::JTT_X86_64_ABS && !getRelocationAt(EntryAddress)) {
         LLVM_DEBUG(
             dbgs()
             << "FAIL: JTT_NORMAL table, no relocation for this address\n");
         break;
       }
     }
 
-    const uint64_t Value =
-        (Type == JumpTable::JTT_PIC)
-            ? Address + *getSignedValueAtAddress(EntryAddress, EntrySize)
-            : *getPointerAtAddress(EntryAddress);
+    uint64_t Value = 0;
+    switch (Type) {
+    case JumpTable::JTT_X86_64_PIC4:
+      Value = Address + *getSignedValueAtAddress(EntryAddress, EntrySize);
+      break;
+    case JumpTable::JTT_X86_64_ABS:
+      Value = *getPointerAtAddress(EntryAddress);
+      break;
+    case JumpTable::JTT_AARCH64_REL1:
+    case JumpTable::JTT_AARCH64_REL2:
+    case JumpTable::JTT_AARCH64_REL4:
+      unsigned ShiftAmt = Type == JumpTable::JTT_AARCH64_REL4 ? 0 : 2;
+      assert(JT &&
+             "jump table must be non-null for AArch64 in analyzeJumpTable");
+      Value = JT->BaseAddress +
+              (*getUnsignedValueAtAddress(EntryAddress, EntrySize) << ShiftAmt);
+      break;
+    }
 
     // __builtin_unreachable() case.
     if (Value == UnreachableAddress) {
@@ -645,24 +655,19 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
 
     // Function or one of its fragments.
     const BinaryFunction *TargetBF = getBinaryFunctionContainingAddress(Value);
-    const bool DoesBelongToFunction =
-        BF.containsAddress(Value) ||
-        (TargetBF && areRelatedFragments(TargetBF, &BF));
-    if (!DoesBelongToFunction) {
+    if (!TargetBF || !areRelatedFragments(TargetBF, &BF)) {
       LLVM_DEBUG({
-        if (!BF.containsAddress(Value)) {
-          dbgs() << "FAIL: function doesn't contain this address\n";
-          if (TargetBF) {
-            dbgs() << "  ! function containing this address: "
-                   << TargetBF->getPrintName() << '\n';
-            if (TargetBF->isFragment()) {
-              dbgs() << "  ! is a fragment";
-              for (BinaryFunction *Parent : TargetBF->ParentFragments)
-                dbgs() << ", parent: " << Parent->getPrintName();
-              dbgs() << '\n';
-            }
-          }
-        }
+        dbgs() << "FAIL: function doesn't contain this address\n";
+        if (!TargetBF)
+          break;
+        dbgs() << "  ! function containing this address: " << *TargetBF << '\n';
+        if (!TargetBF->isFragment())
+          break;
+        dbgs() << "  ! is a fragment with parents: ";
+        ListSeparator LS;
+        for (BinaryFunction *Parent : TargetBF->ParentFragments)
+          dbgs() << LS << *Parent;
+        dbgs() << '\n';
       });
       break;
     }
@@ -677,17 +682,17 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
     ++NumRealEntries;
     LLVM_DEBUG(dbgs() << formatv("OK: {0:x} real entry\n", Value));
 
-    if (TargetBF != &BF && HasEntryInFragment)
-      *HasEntryInFragment = true;
+    if (TargetBF != &BF && JT)
+      JT->IsSplit = true;
     addEntryAddress(Value);
   }
 
   // Trim direct/normal jump table to exclude trailing unreachable entries that
   // can collide with a function address.
-  if (Type == JumpTable::JTT_NORMAL && EntriesAsAddress &&
-      TrimmedSize != EntriesAsAddress->size() &&
+  if (Type == JumpTable::JTT_X86_64_ABS && JT &&
+      TrimmedSize != JT->EntriesAsAddress.size() &&
       getBinaryFunctionAtAddress(UnreachableAddress))
-    EntriesAsAddress->resize(TrimmedSize);
+    JT->EntriesAsAddress.resize(TrimmedSize);
 
   // It's a jump table if the number of real entries is more than 1, or there's
   // one real entry and one or more special targets. If there are only multiple
@@ -702,20 +707,20 @@ void BinaryContext::populateJumpTables() {
        ++JTI) {
     JumpTable *JT = JTI->second;
 
-    bool NonSimpleParent = false;
-    for (BinaryFunction *BF : JT->Parents)
-      NonSimpleParent |= !BF->isSimple();
-    if (NonSimpleParent)
+    auto isSimple = std::bind(&BinaryFunction::isSimple, std::placeholders::_1);
+    if (!llvm::all_of(JT->Parents, isSimple))
       continue;
 
     uint64_t NextJTAddress = 0;
     auto NextJTI = std::next(JTI);
-    if (NextJTI != JTE)
+    if (isAArch64()) {
+      NextJTAddress = JT->getAddress() + JT->getSize();
+      JT->Entries.clear();
+    } else if (NextJTI != JTE)
       NextJTAddress = NextJTI->second->getAddress();
 
-    const bool Success =
-        analyzeJumpTable(JT->getAddress(), JT->Type, *(JT->Parents[0]),
-                         NextJTAddress, &JT->EntriesAsAddress, &JT->IsSplit);
+    const bool Success = analyzeJumpTable(
+        JT->getAddress(), JT->Type, *JT->Parents.front(), NextJTAddress, JT);
     if (!Success) {
       LLVM_DEBUG({
         dbgs() << "failed to analyze ";
@@ -743,7 +748,7 @@ void BinaryContext::populateJumpTables() {
 
     // In strict mode, erase PC-relative relocation record. Later we check that
     // all such records are erased and thus have been accounted for.
-    if (opts::StrictMode && JT->Type == JumpTable::JTT_PIC) {
+    if (opts::StrictMode && JT->Type == JumpTable::JTT_X86_64_PIC4) {
       for (uint64_t Address = JT->getAddress();
            Address < JT->getAddress() + JT->getSize();
            Address += JT->EntrySize) {
@@ -839,33 +844,26 @@ BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
     assert(JT->Type == Type && "jump table types have to match");
     assert(Address == JT->getAddress() && "unexpected non-empty jump table");
 
-    // Prevent associating a jump table to a specific fragment twice.
-    if (!llvm::is_contained(JT->Parents, &Function)) {
-      assert(llvm::all_of(JT->Parents,
-                          [&](const BinaryFunction *BF) {
-                            return areRelatedFragments(&Function, BF);
-                          }) &&
-             "cannot re-use jump table of a different function");
-      // Duplicate the entry for the parent function for easy access
-      JT->Parents.push_back(&Function);
-      if (opts::Verbosity > 2) {
-        this->outs() << "BOLT-INFO: Multiple fragments access same jump table: "
-                     << JT->Parents[0]->getPrintName() << "; "
-                     << Function.getPrintName() << "\n";
-        JT->print(this->outs());
-      }
-      Function.JumpTables.emplace(Address, JT);
-      for (BinaryFunction *Parent : JT->Parents)
-        Parent->setHasIndirectTargetToSplitFragment(true);
-    }
+    if (llvm::is_contained(JT->Parents, &Function))
+      return JT->getFirstLabel();
 
-    bool IsJumpTableParent = false;
-    (void)IsJumpTableParent;
-    for (BinaryFunction *Frag : JT->Parents)
-      if (Frag == &Function)
-        IsJumpTableParent = true;
-    assert(IsJumpTableParent &&
+    // Prevent associating a jump table to a specific fragment twice.
+    auto isSibling = std::bind(&BinaryContext::areRelatedFragments, this,
+                               &Function, std::placeholders::_1);
+    assert(llvm::all_of(JT->Parents, isSibling) &&
            "cannot re-use jump table of a different function");
+    if (opts::Verbosity > 2) {
+      this->outs() << "BOLT-INFO: Multiple fragments access same jump table: "
+                   << JT->Parents[0]->getPrintName() << "; "
+                   << Function.getPrintName() << "\n";
+      JT->print(this->outs());
+    }
+    if (JT->Parents.size() == 1)
+      JT->Parents.front()->setHasIndirectTargetToSplitFragment(true);
+    Function.setHasIndirectTargetToSplitFragment(true);
+    // Duplicate the entry for the parent function for easy access
+    JT->Parents.push_back(&Function);
+    Function.JumpTables.emplace(Address, JT);
     return JT->getFirstLabel();
   }