[BOLT][AArch64] Add jump table support using .llvm_jump_table_info

bolt/include/bolt/Core/JumpTable.h

@@ -40,6 +40,7 @@ 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;
@@ -116,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/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/lib/Core/BinaryContext.cpp

@@ -625,6 +625,15 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
     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.
@@ -704,7 +713,10 @@ void BinaryContext::populateJumpTables() {
 
     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(

bolt/lib/Rewrite/CMakeLists.txt

@@ -17,6 +17,7 @@ add_llvm_library(LLVMBOLTRewrite
   DWARFRewriter.cpp
   ExecutableFileMemoryManager.cpp
   JITLinkLinker.cpp
+  JumpTableInfoReader.cpp
   LinuxKernelRewriter.cpp
   MachORewriteInstance.cpp
   MetadataManager.cpp

bolt/lib/Rewrite/JumpTableInfoReader.cpp

@@ -0,0 +1,93 @@
+//===- bolt/Rewrite/JumpTableInfoReader.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Read .llvm_jump_table_info section and register jump tables.
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Core/JumpTable.h"
+#include "bolt/Rewrite/MetadataRewriter.h"
+#include "bolt/Rewrite/MetadataRewriters.h"
+#include "llvm/Support/DataExtractor.h"
+
+using namespace llvm;
+using namespace bolt;
+
+namespace {
+class JumpTableInfoReader final : public MetadataRewriter {
+
+public:
+  JumpTableInfoReader(StringRef Name, BinaryContext &BC)
+      : MetadataRewriter(Name, BC) {}
+  Error preDisasmInitializer() override;
+};
+
+Error JumpTableInfoReader::preDisasmInitializer() {
+  if (!BC.isAArch64())
+    return Error::success();
+
+  ErrorOr<BinarySection &> ErrorOrJTInfoSection =
+      BC.getUniqueSectionByName(".llvm_jump_table_info");
+  if (std::error_code E = ErrorOrJTInfoSection.getError())
+    return Error::success();
+  BinarySection &JTInfoSection = *ErrorOrJTInfoSection;
+  StringRef Buf = JTInfoSection.getContents();
+  DataExtractor DE = DataExtractor(Buf, BC.AsmInfo->isLittleEndian(),
+                                   BC.AsmInfo->getCodePointerSize());
+  DataExtractor::Cursor Cursor(0);
+  while (Cursor && !DE.eof(Cursor)) {
+    const uint8_t Format = DE.getU8(Cursor);
+    const uint64_t JTAddr = DE.getAddress(Cursor);
+    const uint64_t JTBase = DE.getAddress(Cursor);
+    const uint64_t JTLoad = DE.getAddress(Cursor);
+    const uint64_t Branch = DE.getAddress(Cursor);
+    const uint64_t NumEntries = DE.getULEB128(Cursor);
+
+    JumpTable::JumpTableType Type = JumpTable::JTT_AARCH64_LAST;
+    switch (Format) {
+    case 2:
+      Type = JumpTable::JTT_AARCH64_REL1;
+      break;
+    case 3:
+      Type = JumpTable::JTT_AARCH64_REL2;
+      break;
+    case 4:
+      Type = JumpTable::JTT_AARCH64_REL4;
+      break;
+    }
+
+    if (Type == JumpTable::JTT_AARCH64_LAST) {
+      errs() << "BOLT-WARNING: unknown jump table info type " << Format
+                << " for jump table " << Twine::utohexstr(JTAddr) << '\n';
+      continue;
+    }
+
+    BinaryFunction *BF = BC.getBinaryFunctionContainingAddress(Branch);
+    if (!BF) {
+      BC.errs() << "BOLT-WARNING: binary function not found for jump table "
+                   "with address "
+                << Twine::utohexstr(JTAddr) << " and branch "
+                << Twine::utohexstr(Branch) << '\n';
+      continue;
+    }
+    const MCSymbol *JTSym = BC.getOrCreateJumpTable(*BF, JTAddr, Type);
+    assert(JTSym && "failed to create a jump table");
+    JumpTable *JT = BC.getJumpTableContainingAddress(JTAddr);
+    assert(JT && "internal error creating jump table");
+    JT->BaseAddress = JTBase;
+    JT->MemLocInstrAddress = JTLoad;
+    JT->Entries.resize(NumEntries);
+  }
+  return Cursor.takeError();
+}
+} // namespace
+
+std::unique_ptr<MetadataRewriter>
+llvm::bolt::createJumpTableInfoReader(BinaryContext &BC) {
+  return std::make_unique<JumpTableInfoReader>("jump-table-info-reader", BC);
+}

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -3227,6 +3227,8 @@ void RewriteInstance::initializeMetadataManager() {
   MetadataManager.registerRewriter(createPseudoProbeRewriter(*BC));
 
   MetadataManager.registerRewriter(createSDTRewriter(*BC));
+
+  MetadataManager.registerRewriter(createJumpTableInfoReader(*BC));
 }
 
 void RewriteInstance::processSectionMetadata() {