[RFC][BOLT] Create an independent namespace for jump tables

[rafauler]

Summary

Make jump tables inaccessible via our global symbol
table. Make jump table symbols only accessible via
BinaryContext::getOrCreateJumpTable(), called by our jump table
recognition heuristics. In practice, this means that any data access
in code will not reference a jump table, unless disassembly/cfg
analysis explictly recognize the code fragment as accessing a jump
table.

This solves bugs related to regular data references in code
accidentally being bound to a jump table, and this reference being
updated to a new (incorrect) location because we moved this jump
table.

An alternative solution to this would be to honor the connection of
JumpTable objects with their respective BinaryData entries (in the
global symbol table). To make this consistent, we would further need
to change jump table emission to emit all labels in BinaryData at that
location whenever emitting a jump table. Currently, this is
inconsistent and jump tables are poorly synced with the global symbol
table. This solution is more complicated than severing the connection
between jump table and the global symbols (do not publish a jump table
in our global symbol table), and will also lead to less functions
being processed (albeit this number should be small).

Fixes #55004

[rafauler]

This diff is published for discussion only. The alternative implementation that is the most popular solution is in https://reviews.llvm.org/D134098.

bolt/include/bolt/Core/BinaryContext.h

public:
 
   /// [start memory address] -> [segment info] mapping.
   std::map<uint64_t, SegmentInfo> SegmentMapInfo;
 
   /// Symbols that are expected to be undefined in MCContext during emission.
   std::unordered_set<MCSymbol *> UndefinedSymbols;
 
   /// [name] -> [BinaryData*] map used for global symbol resolution.
-  using SymbolMapType = StringMap<BinaryData *>;
-  SymbolMapType GlobalSymbols;
+  StringMap<BinaryData *> GlobalSymbols;
+  StringMap<JumpTable *> NameToJumpTable;
 
   /// [address] -> [BinaryData], ...
   /// Addresses never change.
-  /// Note: it is important that clients do not hold on to instances of
-  /// BinaryData* while the map is still being modified during BinaryFunction
-  /// disassembly.  This is because of the possibility that a regular
-  /// BinaryData is later discovered to be a JumpTable.
   using BinaryDataMapType = std::map<uint64_t, BinaryData *>;
   using binary_data_iterator = BinaryDataMapType::iterator;
   using binary_data_const_iterator = BinaryDataMapType::const_iterator;
   BinaryDataMapType BinaryDataMap;
 
   using FilteredBinaryDataConstIterator =
       FilterIterator<binary_data_const_iterator>;
   using FilteredBinaryDataIterator = FilterIterator<binary_data_iterator>;
JumpTable *getJumpTableContainingAddress(uint64_t Address) {
     --JTI;
     if (JTI->first + JTI->second->getSize() > Address)
       return JTI->second;
     if (JTI->second->getSize() == 0 && JTI->first == Address)
       return JTI->second;
     return nullptr;
   }
 
+  /// Return JumpTable starting at a specific \p Address.
+  JumpTable *getJumpTableAtAddress(uint64_t Address) {
+    auto JTI = JumpTables.find(Address);
+    if (JTI != JumpTables.end())
+      return JTI->second;
+    return nullptr;
+  }
+
   unsigned getDWARFEncodingSize(unsigned Encoding) {
     switch (Encoding & 0x0f) {
     default:
       llvm_unreachable("unknown encoding");
     case dwarf::DW_EH_PE_absptr:
     case dwarf::DW_EH_PE_signed:
       return AsmInfo->getCodePointerSize();
     case dwarf::DW_EH_PE_udata2:
public:
   /// Analyze memory contents at the given \p Address and return the type of
   /// memory contents (such as a possible jump table).
   MemoryContentsType analyzeMemoryAt(uint64_t Address, BinaryFunction &BF);
 
   /// Return a value of the global \p Symbol or an error if the value
   /// was not set.
   ErrorOr<uint64_t> getSymbolValue(const MCSymbol &Symbol) const {
     const BinaryData *BD = getBinaryDataByName(Symbol.getName());
-    if (!BD)
+    if (!BD) {
+      if (const JumpTable *JT = getJumpTableByName(Symbol.getName()))
+        return JT->getAddress();
       return std::make_error_code(std::errc::bad_address);
+    }
     return BD->getAddress();
   }
 
+  const JumpTable *getJumpTableByName(StringRef Name) const {
+    auto Itr = NameToJumpTable.find(Name);
+    return Itr != NameToJumpTable.end() ? Itr->second : nullptr;
+  }
+
   /// Return a global symbol registered at a given \p Address and \p Size.
   /// If no symbol exists, create one with unique name using \p Prefix.
   /// If there are multiple symbols registered at the \p Address, then
   /// return the first one.
   MCSymbol *getOrCreateGlobalSymbol(uint64_t Address, Twine Prefix,
                                     uint64_t Size = 0, uint16_t Alignment = 0,
                                     unsigned Flags = 0);
 

bolt/lib/Core/BinaryContext.cpp

if (JumpTable *JT = getJumpTableContainingAddress(Address)) {
     for (BinaryFunction *Frag : JT->Parents)
       if (Frag == &Function)
         IsJumpTableParent = true;
     assert(IsJumpTableParent &&
            "cannot re-use jump table of a different function");
     return JT->getFirstLabel();
   }
 
-  // Re-use the existing symbol if possible.
-  MCSymbol *JTLabel = nullptr;
-  if (BinaryData *Object = getBinaryDataAtAddress(Address)) {
-    if (!isInternalSymbolName(Object->getSymbol()->getName()))
-      JTLabel = Object->getSymbol();
-  }
-
   const uint64_t EntrySize = getJumpTableEntrySize(Type);
-  if (!JTLabel) {
-    const std::string JumpTableName = generateJumpTableName(Function, Address);
-    JTLabel = registerNameAtAddress(JumpTableName, Address, 0, EntrySize);
-  }
+  MCSymbol *JTLabel = Ctx->getOrCreateSymbol(JumpTableName);
 
   LLVM_DEBUG(dbgs() << "BOLT-DEBUG: creating jump table " << JTLabel->getName()
                     << " in function " << Function << '\n');
 
   JumpTable *JT = new JumpTable(*JTLabel, Address, EntrySize, Type,
                                 JumpTable::LabelMapType{{0, JTLabel}},
                                 *getSectionForAddress(Address));
   JT->Parents.push_back(&Function);
   if (opts::Verbosity > 2)
     JT->print(outs());
   JumpTables.emplace(Address, JT);
+  NameToJumpTable[JTLabel->getName()] = JT;
 
   // Duplicate the entry for the parent function for easy access.
   Function.JumpTables.emplace(Address, JT);
   return JTLabel;
 }
 
 std::pair<uint64_t, const MCSymbol *>
 BinaryContext::duplicateJumpTable(BinaryFunction &Function, JumpTable *JT,
BinaryContext::duplicateJumpTable(BinaryFunction &Function, JumpTable *JT,
   NewJT->Parents = JT->Parents;
   NewJT->Entries = JT->Entries;
   NewJT->Counts = JT->Counts;
   uint64_t JumpTableID = ++DuplicatedJumpTables;
   // Invert it to differentiate from regular jump tables whose IDs are their
   // addresses in the input binary memory space
   JumpTableID = ~JumpTableID;
   JumpTables.emplace(JumpTableID, NewJT);
+  NameToJumpTable[NewLabel->getName()] = NewJT;
   Function.JumpTables.emplace(JumpTableID, NewJT);
   return std::make_pair(JumpTableID, NewLabel);
 }
 
 std::string BinaryContext::generateJumpTableName(const BinaryFunction &BF,
                                                  uint64_t Address) {
   size_t Id;
   uint64_t Offset = 0;

bolt/lib/Core/BinaryFunction.cpp

if (JT.Entries.empty()) {
               (HasOneParent && TargetBF == this)
                   ? getOrCreateLocalLabel(JT.EntriesAsAddress[I], true)
                   : TargetBF->addEntryPointAtOffset(EntryAddress -
                                                     TargetBF->getAddress());
           JT.Entries.push_back(Label);
         }
       }
     }
-
-    const uint64_t BDSize =
-        BC.getBinaryDataAtAddress(JT.getAddress())->getSize();
-    if (!BDSize) {
-      BC.setBinaryDataSize(JT.getAddress(), JT.getSize());
-    } else {
-      assert(BDSize >= JT.getSize() &&
-             "jump table cannot be larger than the containing object");
-    }
   }
 
   // Add TakenBranches from JumpTables.
   //
   // We want to do it after initial processing since we don't know jump tables'
   // boundaries until we process them all.
   for (auto &JTSite : JTSites) {
     const uint64_t JTSiteOffset = JTSite.first;

bolt/lib/Passes/IdenticalCodeFolding.cpp

while (I != E && OtherI != OtherE) {
           return FunctionA == FunctionB;
         }
 
         // One of the symbols represents a function, the other one does not.
         if (FunctionA != FunctionB)
           return false;
 
         // Check if symbols are jump tables.
-        const BinaryData *SIA = BC.getBinaryDataByName(SymbolA->getName());
-        if (!SIA)
-          return false;
-        const BinaryData *SIB = BC.getBinaryDataByName(SymbolB->getName());
-        if (!SIB)
-          return false;
-
-        assert((SIA->getAddress() != SIB->getAddress()) &&
-               "different symbols should not have the same value");
-
-        const JumpTable *JumpTableA =
-            A.getJumpTableContainingAddress(SIA->getAddress());
+        const JumpTable *JumpTableA = BC.getJumpTableByName(SymbolA->getName());
         if (!JumpTableA)
           return false;
 
-        const JumpTable *JumpTableB =
-            B.getJumpTableContainingAddress(SIB->getAddress());
+        const JumpTable *JumpTableB = BC.getJumpTableByName(SymbolB->getName());
         if (!JumpTableB)
           return false;
 
-        if ((SIA->getAddress() - JumpTableA->getAddress()) !=
-            (SIB->getAddress() - JumpTableB->getAddress()))
-          return false;
-
         return equalJumpTables(*JumpTableA, *JumpTableB, A, B);
       };
 
       if (!isInstrEquivalentWith(*I, *BB, *OtherI, *OtherBB,
                                  AreSymbolsIdentical))
         return false;
 
       ++I;

bolt/lib/Profile/DataAggregator.cpp

if (!Func) {
       });
       continue;
     }
 
     FuncName = Func->getOneName();
     PC -= Func->getAddress();
 
     // Try to resolve symbol for memory load
-    if (BinaryData *BD = BC->getBinaryDataContainingAddress(Addr)) {
+    if (JumpTable *JT = BC->getJumpTableContainingAddress(Addr)) {
+      MemName = JT->getName();
+      Addr -= JT->getAddress();
+    } else if (BinaryData *BD = BC->getBinaryDataContainingAddress(Addr)) {
       MemName = BD->getName();
       Addr -= BD->getAddress();
     } else if (opts::FilterMemProfile) {
       // Filter out heap/stack accesses
       continue;
     }
 
     const Location FuncLoc(!FuncName.empty(), FuncName, PC);

bolt/lib/Profile/DataReader.cpp

for (MemInfo &MI : MemoryData->Data) {
         // Ignore bad instruction address.
         continue;
       }
 
       auto &MemAccessProfile =
           BC.MIB->getOrCreateAnnotationAs<MemoryAccessProfile>(
               II->second, "MemoryAccessProfile");
       BinaryData *BD = nullptr;
-      if (MI.Addr.IsSymbol)
+      if (MI.Addr.IsSymbol) {
+        if (const JumpTable *JT = BC.getJumpTableByName(MI.Addr.Name))
+          BD = BC.getJumpTableAtAddress(JT->getAddress());
+        else
           BD = BC.getBinaryDataByName(MI.Addr.Name);
+      }
       MemAccessProfile.AddressAccessInfo.push_back(
           {BD, MI.Addr.Offset, MI.Count});
       auto NextII = std::next(II);
       if (NextII == Function.Instructions.end())
         MemAccessProfile.NextInstrOffset = Function.getSize();
       else
         MemAccessProfile.NextInstrOffset = II->first;
     }

bolt/lib/Rewrite/RewriteInstance.cpp

for (const StringRef &Symbol : Symbols) {
         uint64_t Address = I->isMoved() && !I->isJumpTable()
                                ? I->getOutputAddress()
                                : I->getAddress();
         LLVM_DEBUG(dbgs() << "Resolved to address 0x"
                           << Twine::utohexstr(Address) << "\n");
         AllResults[Symbol] = JITEvaluatedSymbol(Address, JITSymbolFlags());
         continue;
       }
+      if (const JumpTable *JT = BC.getJumpTableByName(SymName)) {
+        AllResults[Symbol] = JITEvaluatedSymbol(JT->getAddress(), JITSymbolFlags());
+        continue;
+      }
       LLVM_DEBUG(dbgs() << "Resolved to address 0x0\n");
       AllResults[Symbol] = JITEvaluatedSymbol(0, JITSymbolFlags());
     }
 
     OnResolved(std::move(AllResults));
   }
 };
 

bolt/test/X86/split-func-jump-table-fragment-bidirection.s

 # This reproduces an issue where two fragments of same function access same
 # jump table, which means at least one fragment visits the other, i.e., one
 # of them has split jump table. As a result, all of them will be marked as
 # non-simple function.
 
 # REQUIRES: system-linux
 
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
 # RUN: llvm-strip --strip-unneeded %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
 # RUN: llvm-bolt -print-cfg -v=3 %t.exe -o %t.out 2>&1 | FileCheck %s
 
 # CHECK: BOLT-INFO: Multiple fragments access same jump table: main; main.cold.1
-# CHECK: PIC Jump table JUMP_TABLE1 for function main, main.cold.1 at {{.*}}  with a total count of 0:
+# CHECK: PIC Jump table {{.*}} for function main, main.cold.1 at {{.*}}  with a total count of 0:
 
   .text
   .globl main
   .type main, %function
   .p2align 2
 main:
 LBB0:
   andl $0xf, %ecx

bolt/test/runtime/X86/jt-symbol-disambiguation.s

+# In this test case, the symbol that represents the end of a table
+# in .rodata is being colocated with the start of a jump table from
+# another function, and BOLT moves that jump table. This should not
+# cause the symbol representing the end of the table to be moved as
+# well.
+# Bug reported in https://github.com/llvm/llvm-project/issues/55004
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags -no-pie -nostartfiles -nostdlib -lc %t.o -o %t.exe -Wl,-q
+
+# RUN: llvm-bolt %t.exe -o %t.exe.bolt --relocs=1 --lite=0 \
+# RUN:   --reorder-blocks=reverse -jump-tables=move
+
+# RUN: %t.exe.bolt 1 2 3
+
+  .file "jt-symbol-disambiguation.s"
+  .text
+
+# ----
+# Func foo contains a jump table whose start is colocated with a
+# symbol marking the end of a data table
+# ----
+  .globl foo
+  .type foo, @function
+foo:
+  .cfi_startproc
+  xor    %rax,%rax
+  and    $0x3,%rdi
+  leaq   .JT1(%rip), %rax
+  movslq  (%rax, %rdi, 4), %rdi
+  addq   %rax, %rdi
+  jmpq   *%rdi
+.LBB1:
+  movl   $0x1,%eax
+  jmp    .LBB5
+.LBB2:
+  movl   $0x2,%eax
+  jmp    .LBB5
+.LBB3:
+  movl   $0x3,%eax
+  jmp    .LBB5
+.LBB4:
+  movl   $0x4,%eax
+.LBB5:
+  retq
+  .cfi_endproc
+  .size foo, .-foo
+
+# ----
+# Func _start scans a table using begin/end pointers. End pointer is colocated
+# with the start of a jump table of function foo. When that jump
+# table moves, end pointer in _start should not be affected.
+# ----
+  .globl _start
+  .type _start, @function
+_start:
+  .cfi_startproc
+  movq   (%rsp), %rdi
+  callq foo
+  leaq   .start_of_table(%rip), %rsi  # iterator
+  leaq   .end_of_table(%rip), %rdi    # iterator end
+.LBB6:
+  cmpq %rsi, %rdi
+  je .LBB7
+  movq (%rsi), %rbx
+  leaq 8(%rsi), %rsi            # ++iterator
+  jmp .LBB6
+.LBB7:
+  xor   %rdi, %rdi
+  callq exit@PLT
+  .cfi_endproc
+  .size _start, .-_start
+
+# ----
+# Data section
+# ----
+  .section .rodata,"a",@progbits
+  .p2align 3
+.start_of_table:
+  .quad 123
+  .quad 456
+  .quad 789
+.end_of_table:
+.JT1:
+  .long .LBB1 - .JT1
+  .long .LBB2 - .JT1
+  .long .LBB3 - .JT1
+  .long .LBB4 - .JT1