[BOLT] Move JumpTable management to BinaryContext

Summary:
Make BinaryContext responsible for creation and management of
JumpTables. This will be used for detection and resolution of jump table
conflicts across functions.

(cherry picked from FBD15196017)

bolt/src/BinaryContext.cpp

@@ -253,6 +253,81 @@ BinaryFunction *BinaryContext::createBinaryFunction(
   return BF;
 }
 
+std::pair<JumpTable *, const MCSymbol *>
+BinaryContext::createJumpTable(BinaryFunction &Function,
+                               uint64_t Address,
+                               JumpTable::JumpTableType Type,
+                               JumpTable::OffsetEntriesType &&OffsetEntries) {
+  const auto JumpTableName = generateJumpTableName(Function, Address);
+  if (auto *JT = getJumpTableContainingAddress(Address)) {
+    assert(JT->Type == Type && "jump table types have to match");
+    assert(JT->Parent == &Function &&
+           "cannot re-use jump table of a different function");
+    assert((Address == JT->getAddress() || Type != JumpTable::JTT_PIC) &&
+           "cannot re-use part of PIC jump table");
+    // Get or create a new label for the table.
+    const auto JTOffset = Address - JT->getAddress();
+    auto LI = JT->Labels.find(JTOffset);
+    if (LI == JT->Labels.end()) {
+      auto *JTStartLabel = registerNameAtAddress(JumpTableName,
+                                                 Address,
+                                                 0,
+                                                 JT->EntrySize);
+      auto Result = JT->Labels.emplace(JTOffset, JTStartLabel);
+      assert(Result.second && "error adding jump table label");
+      LI = Result.first;
+    }
+
+    return std::make_pair(JT, LI->second);
+  }
+
+  auto *JTStartLabel = Ctx->getOrCreateSymbol(JumpTableName);
+  const auto EntrySize =
+    Type == JumpTable::JTT_PIC ? 4 : AsmInfo->getCodePointerSize();
+
+  DEBUG(dbgs() << "BOLT-DEBUG: creating jump table "
+               << JTStartLabel->getName()
+               << " in function " << Function << " with "
+               << OffsetEntries.size() << " entries\n");
+
+  auto *JT = new JumpTable(JumpTableName,
+                           Address,
+                           EntrySize,
+                           Type,
+                           std::move(OffsetEntries),
+                           JumpTable::LabelMapType{{0, JTStartLabel}},
+                           Function,
+                           *getSectionForAddress(Address));
+
+  const auto *JTLabel = registerNameAtAddress(JumpTableName, Address, JT);
+  assert(JTLabel == JTStartLabel);
+
+  JumpTables.emplace(Address, JT);
+
+  // Duplicate the entry for the parent function for easy access.
+  Function.JumpTables.emplace(Address, JT);
+
+  return std::make_pair(JT, JTLabel);
+}
+
+std::string BinaryContext::generateJumpTableName(const BinaryFunction &BF,
+                                                 uint64_t Address) {
+  size_t Id;
+  uint64_t Offset = 0;
+  if (const auto *JT = BF.getJumpTableContainingAddress(Address)) {
+    Offset = Address - JT->getAddress();
+    auto Itr = JT->Labels.find(Offset);
+    if (Itr != JT->Labels.end()) {
+      return Itr->second->getName();
+    }
+    Id = JumpTableIds.at(JT->getAddress());
+  } else {
+    Id = JumpTableIds[Address] = BF.JumpTables.size();
+  }
+  return ("JUMP_TABLE/" + BF.Names[0] + "." + std::to_string(Id) +
+          (Offset ? ("." + std::to_string(Offset)) : ""));
+}
+
 MCSymbol *BinaryContext::registerNameAtAddress(StringRef Name,
                                                uint64_t Address,
                                                uint64_t Size,

bolt/src/BinaryContext.h

@@ -17,6 +17,7 @@
 #include "BinaryData.h"
 #include "BinarySection.h"
 #include "DebugData.h"
+#include "JumpTable.h"
 #include "MCPlusBuilder.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/Triple.h"
@@ -145,6 +146,9 @@ class BinaryContext {
   /// Functions injected by BOLT
   std::vector<BinaryFunction *> InjectedBinaryFunctions;
 
+  /// Jump tables for all functions mapped by address.
+  std::map<uint64_t, JumpTable *> JumpTables;
+
 public:
   /// [name] -> [BinaryData*] map used for global symbol resolution.
   using SymbolMapType = std::map<std::string, BinaryData *>;
@@ -199,6 +203,18 @@ class BinaryContext {
         getBinaryFunctionAtAddress(Address, Shallow);
   }
 
+  /// Return JumpTable containing a given \p Address.
+  JumpTable *getJumpTableContainingAddress(uint64_t Address) {
+    auto JTI = JumpTables.upper_bound(Address);
+    if (JTI == JumpTables.begin())
+      return nullptr;
+    --JTI;
+    if (JTI->first + JTI->second->getSize() > Address) {
+      return JTI->second;
+    }
+    return nullptr;
+  }
+
   /// [MCSymbol] -> [BinaryFunction]
   ///
   /// As we fold identical functions, multiple symbols can point
@@ -272,6 +288,19 @@ class BinaryContext {
     return InjectedBinaryFunctions;
   }
 
+  /// Construct a jump table for \p Function at \p Address.
+  /// May create an embedded jump table and return its label as the second
+  /// element of the pair.
+  std::pair<JumpTable *, const MCSymbol *>
+  createJumpTable(BinaryFunction &Function,
+                  uint64_t Address,
+                  JumpTable::JumpTableType Type,
+                  JumpTable::OffsetEntriesType &&OffsetEntries);
+
+  /// Generate a unique name for jump table at a given \p Address belonging
+  /// to function \p BF.
+  std::string generateJumpTableName(const BinaryFunction &BF, uint64_t Address);
+
 public:
   /// Regular page size.
   static constexpr unsigned RegularPageSize = 0x1000;
@@ -282,6 +311,10 @@ class BinaryContext {
   /// Map address to a constant island owner (constant data in code section)
   std::map<uint64_t, BinaryFunction *> AddressToConstantIslandMap;
 
+  /// A map from jump table address to insertion order.  Used for generating
+  /// jump table names.
+  std::map<uint64_t, size_t> JumpTableIds;
+
   /// Set of addresses in the code that are not a function start, and are
   /// referenced from outside of containing function. E.g. this could happen
   /// when a function has more than a single entry point.

bolt/src/BinaryFunction.cpp

@@ -772,9 +772,27 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction,
   DEBUG(dbgs() << "BOLT-DEBUG: addressed memory is 0x"
                << Twine::utohexstr(ArrayStart) << '\n');
 
+  // List of possible jump targets.
+  std::vector<uint64_t> JTOffsetCandidates;
+
+  auto useJumpTableForInstruction = [&](JumpTable::JumpTableType JTType) {
+    JumpTable *JT;
+    const MCSymbol *JTLabel;
+    std::tie(JT, JTLabel) = BC.createJumpTable(*this,
+                                               ArrayStart,
+                                               JTType,
+                                               std::move(JTOffsetCandidates));
+
+    BC.MIB->replaceMemOperandDisp(const_cast<MCInst &>(*MemLocInstr),
+                                  JTLabel, BC.Ctx.get());
+    BC.MIB->setJumpTable(Instruction, ArrayStart, IndexRegNum);
+
+    JTSites.emplace_back(Offset, ArrayStart);
+  };
+
   // Check if there's already a jump table registered at this address.
-  if (auto *JT = getJumpTableContainingAddress(ArrayStart)) {
-    auto JTOffset = ArrayStart - JT->getAddress();
+  if (auto *JT = BC.getJumpTableContainingAddress(ArrayStart)) {
+    const auto JTOffset = ArrayStart - JT->getAddress();
     if (Type == IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE && JTOffset != 0) {
         // Adjust the size of this jump table and create a new one if necessary.
         // We cannot re-use the entries since the offsets are relative to the
@@ -783,7 +801,7 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction,
                      << Twine::utohexstr(JT->getAddress()) << '\n');
         JT->OffsetEntries.resize(JTOffset / JT->EntrySize);
     } else if (Type != IndirectBranchType::POSSIBLE_FIXED_BRANCH) {
-      // Re-use an existing jump table. Perhaps parts of it.
+      // Re-use the existing jump table or parts of it.
       if (Type != IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE) {
         assert(JT->Type == JumpTable::JTT_NORMAL &&
                "normal jump table expected");
@@ -792,24 +810,7 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction,
         assert(JT->Type == JumpTable::JTT_PIC && "PIC jump table expected");
       }
 
-      // Get or create a new label for the table.
-      auto LI = JT->Labels.find(JTOffset);
-      if (LI == JT->Labels.end()) {
-        auto *JTStartLabel =
-          BC.registerNameAtAddress(generateJumpTableName(ArrayStart),
-                                   ArrayStart,
-                                   0,
-                                   JT->EntrySize);
-        auto Result = JT->Labels.emplace(JTOffset, JTStartLabel);
-        assert(Result.second && "error adding jump table label");
-        LI = Result.first;
-      }
-
-      BC.MIB->replaceMemOperandDisp(const_cast<MCInst &>(*MemLocInstr),
-                                    LI->second, BC.Ctx.get());
-      BC.MIB->setJumpTable(Instruction, ArrayStart, IndexRegNum);
-
-      JTSites.emplace_back(Offset, ArrayStart);
+      useJumpTableForInstruction(JT->Type);
 
       return Type;
     }
@@ -839,7 +840,6 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction,
   DataExtractor DE(SectionContents, BC.AsmInfo->isLittleEndian(), EntrySize);
   auto ValueOffset = static_cast<uint32_t>(ArrayStart - Section->getAddress());
   uint64_t Value = 0;
-  std::vector<uint64_t> JTOffsetCandidates;
   auto UpperBound = Section->getSize();
   const auto *JumpTableBD = BC.getBinaryDataAtAddress(ArrayStart);
   if (JumpTableBD && JumpTableBD->getSize()) {
@@ -897,37 +897,10 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction,
     assert(JTOffsetCandidates.size() > 1 &&
            "expected more than one jump table entry");
 
-    auto JumpTableName = generateJumpTableName(ArrayStart);
-    auto JumpTableType =
-      Type == IndirectBranchType::POSSIBLE_JUMP_TABLE
+    const auto JumpTableType = Type == IndirectBranchType::POSSIBLE_JUMP_TABLE
         ? JumpTable::JTT_NORMAL
         : JumpTable::JTT_PIC;
-
-    auto *JTStartLabel = BC.Ctx->getOrCreateSymbol(JumpTableName);
-
-    auto JT = llvm::make_unique<JumpTable>(JumpTableName,
-                                           ArrayStart,
-                                           EntrySize,
-                                           JumpTableType,
-                                           std::move(JTOffsetCandidates),
-                                           JumpTable::LabelMapType{{0, JTStartLabel}},
-                                           *BC.getSectionForAddress(ArrayStart));
-
-    auto *JTLabel = BC.registerNameAtAddress(JumpTableName,
-                                             ArrayStart,
-                                             JT.get());
-    assert(JTLabel == JTStartLabel);
-
-    DEBUG(dbgs() << "BOLT-DEBUG: creating jump table "
-                 << JTStartLabel->getName()
-                 << " in function " << *this << " with "
-                 << JTOffsetCandidates.size() << " entries.\n");
-    JumpTables.emplace(ArrayStart, JT.release());
-    BC.MIB->replaceMemOperandDisp(const_cast<MCInst &>(*MemLocInstr),
-                                  JTStartLabel, BC.Ctx.get());
-    BC.MIB->setJumpTable(Instruction, ArrayStart, IndexRegNum);
-
-    JTSites.emplace_back(Offset, ArrayStart);
+    useJumpTableForInstruction(JumpTableType);
 
     return Type;
   }
@@ -1558,6 +1531,7 @@ void BinaryFunction::postProcessJumpTables() {
         break;
     }
   }
+  clearList(JTSites);
 
   // Free memory used by jump table offsets.
   for (auto &JTI : JumpTables) {
@@ -3484,24 +3458,6 @@ BinaryFunction::BasicBlockOrderType BinaryFunction::dfs() const {
   return DFS;
 }
 
-std::string BinaryFunction::generateJumpTableName(uint64_t Address) const {
-  auto *JT = getJumpTableContainingAddress(Address);
-  size_t Id;
-  uint64_t Offset = 0;
-  if (JT) {
-    Offset = Address - JT->getAddress();
-    auto Itr = JT->Labels.find(Offset);
-    if (Itr != JT->Labels.end()) {
-      return Itr->second->getName();
-    }
-    Id = JumpTableIds.at(JT->getAddress());
-  } else {
-    Id = JumpTableIds[Address] = JumpTables.size();
-  }
-  return ("JUMP_TABLE/" + Names[0] + "." + std::to_string(Id) +
-          (Offset ? ("." + std::to_string(Offset)) : ""));
-}
-
 std::size_t BinaryFunction::hash(bool Recompute, bool UseDFS) const {
   if (size() == 0)
     return 0;

bolt/src/BinaryFunction.h

@@ -451,25 +451,20 @@ class BinaryFunction {
   /// function and that apply before the entry basic block).
   CFIInstrMapType CIEFrameInstructions;
 
-  /// All compound jump tables for this function.
+  /// All compound jump tables for this function. This duplicates what's stored
+  /// in the BinaryContext, but additionally it gives quick access for all
+  /// jump tables used by this function.
+  ///
   /// <OriginalAddress> -> <JumpTable *>
   std::map<uint64_t, JumpTable *> JumpTables;
 
-  /// A map from jump table address to insertion order.  Used for generating
-  /// jump table names.
-  mutable std::map<uint64_t, size_t> JumpTableIds;
-
-  /// Generate a unique name for this jump table at the given address that
-  /// should be repeatable no matter what the start address of the table is.
-  std::string generateJumpTableName(uint64_t Address) const;
-
   /// Iterate over all jump tables associated with this function.
   iterator_range<std::map<uint64_t, JumpTable *>::const_iterator>
   jumpTables() const {
     return make_range(JumpTables.begin(), JumpTables.end());
   }
 
-  /// All jump table sites in the function.
+  /// All jump table sites in the function before CFG is built.
   std::vector<std::pair<uint64_t, uint64_t>> JTSites;
 
   /// List of relocations in this function.

bolt/src/JumpTable.cpp

@@ -27,6 +27,23 @@ extern cl::opt<JumpTableSupportLevel> JumpTables;
 extern cl::opt<unsigned> Verbosity;
 }
 
+JumpTable::JumpTable(StringRef Name,
+                     uint64_t Address,
+                     std::size_t EntrySize,
+                     JumpTableType Type,
+                     OffsetEntriesType &&OffsetEntries,
+                     LabelMapType &&Labels,
+                     BinaryFunction &BF,
+                     BinarySection &Section)
+  : BinaryData(Name, Address, 0, EntrySize, Section),
+    EntrySize(EntrySize),
+    OutputEntrySize(EntrySize),
+    Type(Type),
+    OffsetEntries(OffsetEntries),
+    Labels(Labels),
+    Parent(&BF) {
+}
+
 std::pair<size_t, size_t>
 JumpTable::getEntriesForAddress(const uint64_t Addr) const {
   const uint64_t InstOffset = Addr - getAddress();
@@ -174,18 +191,3 @@ void JumpTable::print(raw_ostream &OS) const {
   }
   OS << "\n\n";
 }
-
-JumpTable::JumpTable(StringRef Name,
-                     uint64_t Address,
-                     std::size_t EntrySize,
-                     JumpTableType Type,
-                     decltype(OffsetEntries) &&OffsetEntries,
-                     decltype(Labels) &&Labels,
-                     BinarySection &Section)
-  : BinaryData(Name, Address, 0, EntrySize, Section),
-    EntrySize(EntrySize),
-    OutputEntrySize(EntrySize),
-    Type(Type),
-    OffsetEntries(OffsetEntries),
-    Labels(Labels)
-{ }