[JITLink] Update splitBlock to support splitting into multiple blocks.

LinkGraph::splitBlock used to take a single split-point to split a Block into
two. In the common case where a block needs to be split repeatedly (e.g. in
eh-frame and compact-unwind sections), iterative calls to splitBlock could
lead to poor performance as symbols and edges are repeatedly shuffled to new
blocks.

This commit updates LinkGraph::splitBlock to take a sequence of split offsets,
allowing a block to be split into an arbitrary number of new blocks. Internally,
Symbols and Edges only need to be moved once (directly to whichever new block
they will be associated with), leading to better performance.

On some large MachO object files in an out of tree project this change improved
the performance of splitBlock by several orders of magnitude.

rdar://135820493

Author: lhames

llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -231,6 +231,12 @@ class Block : public Addressable {
   /// Returns the size of this defined addressable.
   size_t getSize() const { return Size; }
 
+  /// Turns this block into a zero-fill block of the given size.
+  void setZeroFillSize(size_t Size) {
+    Data = nullptr;
+    this->Size = Size;
+  }
+
   /// Returns the address range of this defined addressable.
   orc::ExecutorAddrRange getRange() const {
     return orc::ExecutorAddrRange(getAddress(), getSize());
@@ -1166,17 +1172,22 @@ class LinkGraph {
   /// Cache type for the splitBlock function.
   using SplitBlockCache = std::optional<SmallVector<Symbol *, 8>>;
 
-  /// Splits block B at the given index which must be greater than zero.
-  /// If SplitIndex == B.getSize() then this function is a no-op and returns B.
-  /// If SplitIndex < B.getSize() then this function returns a new block
-  /// covering the range [ 0, SplitIndex ), and B is modified to cover the range
-  /// [ SplitIndex, B.size() ).
+  /// Splits block B into a sequence of smaller blocks.
+  ///
+  /// SplitOffsets should be a sequence of ascending offsets in B. The starting
+  /// offset should be greater than zero, and the final offset less than
+  /// B.getSize() - 1.
+  ///
+  /// The resulting seqeunce of blocks will start with the original block B
+  /// (truncated to end at the first split offset) followed by newly introduced
+  /// blocks starting at the subsequent split points.
   ///
   /// The optional Cache parameter can be used to speed up repeated calls to
-  /// splitBlock for a single block. If the value is None the cache will be
-  /// treated as uninitialized and splitBlock will populate it. Otherwise it
-  /// is assumed to contain the list of Symbols pointing at B, sorted in
-  /// descending order of offset.
+  /// splitBlock for blocks within a single Section. If the value is None then
+  /// the cache will be treated as uninitialized and splitBlock will populate
+  /// it. Otherwise it is assumed to contain the list of Symbols pointing at B,
+  /// sorted in descending order of offset.
+  ///
   ///
   /// Notes:
   ///
@@ -1189,18 +1200,56 @@ class LinkGraph {
   ///    LinkGraph does not have. Clients are responsible for ensuring that
   ///    splitBlock is not used in a way that invalidates edges.
   ///
-  /// 2. The newly introduced block will have a new ordinal which will be
-  ///    higher than any other ordinals in the section. Clients are responsible
-  ///    for re-assigning block ordinals to restore a compatible order if
-  ///    needed.
+  /// 2. The newly introduced blocks will have new ordinals that will be higher
+  ///    than any other ordinals in the section. Clients are responsible for
+  ///    re-assigning block ordinals to restore a compatible order if needed.
   ///
   /// 3. The cache is not automatically updated if new symbols are introduced
   ///    between calls to splitBlock. Any newly introduced symbols may be
   ///    added to the cache manually (descending offset order must be
   ///    preserved), or the cache can be set to None and rebuilt by
   ///    splitBlock on the next call.
-  Block &splitBlock(Block &B, size_t SplitIndex,
-                    SplitBlockCache *Cache = nullptr);
+  template <typename SplitOffsetRange>
+  std::vector<Block *> splitBlock(Block &B, SplitOffsetRange &&SplitOffsets,
+                                  LinkGraph::SplitBlockCache *Cache = nullptr) {
+    std::vector<Block *> Blocks;
+    Blocks.push_back(&B);
+
+    if (std::empty(SplitOffsets))
+      return Blocks;
+
+    // Special case zero-fill:
+    if (B.isZeroFill()) {
+      size_t OrigSize = B.getSize();
+      for (Edge::OffsetT Offset : SplitOffsets) {
+        assert(Offset > 0 && Offset < B.getSize() &&
+               "Split offset must be inside block content");
+        Blocks.back()->setZeroFillSize(
+            Offset - (Blocks.back()->getAddress() - B.getAddress()));
+        Blocks.push_back(&createZeroFillBlock(
+            B.getSection(), B.getSize(), B.getAddress() + Offset,
+            B.getAlignment(),
+            (B.getAlignmentOffset() + Offset) % B.getAlignment()));
+      }
+      Blocks.back()->setZeroFillSize(
+          OrigSize - (Blocks.back()->getAddress() - B.getAddress()));
+      return Blocks;
+    }
+
+    // Handle content blocks. We'll just create the blocks with their starting
+    // address and no content here. The bulk of the work is deferred to
+    // splitBlockImpl.
+    for (Edge::OffsetT Offset : SplitOffsets) {
+      assert(Offset > 0 && Offset < B.getSize() &&
+             "Split offset must be inside block content");
+      Blocks.push_back(&createContentBlock(
+          B.getSection(), ArrayRef<char>(), B.getAddress() + Offset,
+          B.getAlignment(),
+          (B.getAlignmentOffset() + Offset) % B.getAlignment()));
+    }
+
+    return splitBlockImpl(std::move(Blocks), Cache);
+  }
 
   /// Add an external symbol.
   /// Some formats (e.g. ELF) allow Symbols to have sizes. For Symbols whose
@@ -1534,6 +1583,9 @@ class LinkGraph {
   void dump(raw_ostream &OS);
 
 private:
+  std::vector<Block *> splitBlockImpl(std::vector<Block *> Blocks,
+                                      SplitBlockCache *Cache);
+
   // Put the BumpPtrAllocator first so that we don't free any of the underlying
   // memory until the Symbol/Addressable destructors have been run.
   BumpPtrAllocator Allocator;

llvm/lib/ExecutionEngine/JITLink/DWARFRecordSectionSplitter.cpp

@@ -78,12 +78,11 @@ Error DWARFRecordSectionSplitter::processBlock(
       StringRef(B.getContent().data(), B.getContent().size()),
       G.getEndianness());
 
+  std::vector<Edge::OffsetT> SplitOffsets;
   while (true) {
-    uint64_t RecordStartOffset = BlockReader.getOffset();
-
     LLVM_DEBUG({
       dbgs() << "    Processing CFI record at "
-             << formatv("{0:x16}", B.getAddress()) << "\n";
+             << (B.getAddress() + BlockReader.getOffset()) << "\n";
     });
 
     uint32_t Length;
@@ -100,17 +99,16 @@ Error DWARFRecordSectionSplitter::processBlock(
         return Err;
     }
 
-    // If this was the last block then there's nothing to split
-    if (BlockReader.empty()) {
-      LLVM_DEBUG(dbgs() << "      Extracted " << B << "\n");
-      return Error::success();
-    }
+    // If this was the last block then there's nothing more to split
+    if (BlockReader.empty())
+      break;
 
-    uint64_t BlockSize = BlockReader.getOffset() - RecordStartOffset;
-    auto &NewBlock = G.splitBlock(B, BlockSize, &Cache);
-    (void)NewBlock;
-    LLVM_DEBUG(dbgs() << "      Extracted " << NewBlock << "\n");
+    SplitOffsets.push_back(BlockReader.getOffset());
   }
+
+  G.splitBlock(B, SplitOffsets);
+
+  return Error::success();
 }
 
 } // namespace jitlink

llvm/lib/ExecutionEngine/JITLink/JITLink.cpp

@@ -167,84 +167,103 @@ Section::~Section() {
     B->~Block();
 }
 
-Block &LinkGraph::splitBlock(Block &B, size_t SplitIndex,
-                             SplitBlockCache *Cache) {
-
-  assert(SplitIndex > 0 && "splitBlock can not be called with SplitIndex == 0");
-
-  // If the split point covers all of B then just return B.
-  if (SplitIndex == B.getSize())
-    return B;
-
-  assert(SplitIndex < B.getSize() && "SplitIndex out of range");
-
-  // Create the new block covering [ 0, SplitIndex ).
-  auto &NewBlock =
-      B.isZeroFill()
-          ? createZeroFillBlock(B.getSection(), SplitIndex, B.getAddress(),
-                                B.getAlignment(), B.getAlignmentOffset())
-          : createContentBlock(
-                B.getSection(), B.getContent().slice(0, SplitIndex),
-                B.getAddress(), B.getAlignment(), B.getAlignmentOffset());
-
-  // Modify B to cover [ SplitIndex, B.size() ).
-  B.setAddress(B.getAddress() + SplitIndex);
-  B.setContent(B.getContent().slice(SplitIndex));
-  B.setAlignmentOffset((B.getAlignmentOffset() + SplitIndex) %
-                       B.getAlignment());
-
-  // Handle edge transfer/update.
-  {
-    // Copy edges to NewBlock (recording their iterators so that we can remove
-    // them from B), and update of Edges remaining on B.
-    std::vector<Block::edge_iterator> EdgesToRemove;
-    for (auto I = B.edges().begin(); I != B.edges().end();) {
-      if (I->getOffset() < SplitIndex) {
-        NewBlock.addEdge(*I);
-        I = B.removeEdge(I);
-      } else {
-        I->setOffset(I->getOffset() - SplitIndex);
-        ++I;
-      }
-    }
+std::vector<Block *> LinkGraph::splitBlockImpl(std::vector<Block *> Blocks,
+                                               SplitBlockCache *Cache) {
+  assert(!Blocks.empty() && "Blocks must at least contain the original block");
+
+  // Fix up content of all blocks.
+  ArrayRef<char> Content = Blocks.front()->getContent();
+  for (size_t I = 0; I != Blocks.size() - 1; ++I) {
+    Blocks[I]->setContent(
+        Content.slice(Blocks[I]->getAddress() - Blocks[0]->getAddress(),
+                      Blocks[I + 1]->getAddress() - Blocks[I]->getAddress()));
   }
+  Blocks.back()->setContent(
+      Content.slice(Blocks.back()->getAddress() - Blocks[0]->getAddress()));
+  bool IsMutable = Blocks[0]->ContentMutable;
+  for (auto *B : Blocks)
+    B->ContentMutable = IsMutable;
 
-  // Handle symbol transfer/update.
+  // Transfer symbols.
   {
-    // Initialize the symbols cache if necessary.
     SplitBlockCache LocalBlockSymbolsCache;
     if (!Cache)
       Cache = &LocalBlockSymbolsCache;
+
+    // Build cache if required.
     if (*Cache == std::nullopt) {
       *Cache = SplitBlockCache::value_type();
-      for (auto *Sym : B.getSection().symbols())
-        if (&Sym->getBlock() == &B)
-          (*Cache)->push_back(Sym);
 
+      for (auto *Sym : Blocks[0]->getSection().symbols())
+        if (&Sym->getBlock() == Blocks[0])
+          (*Cache)->push_back(Sym);
       llvm::sort(**Cache, [](const Symbol *LHS, const Symbol *RHS) {
-        return LHS->getOffset() > RHS->getOffset();
+        return LHS->getAddress() > RHS->getAddress();
       });
     }
-    auto &BlockSymbols = **Cache;
-
-    // Transfer all symbols with offset less than SplitIndex to NewBlock.
-    while (!BlockSymbols.empty() &&
-           BlockSymbols.back()->getOffset() < SplitIndex) {
-      auto *Sym = BlockSymbols.back();
-      // If the symbol extends beyond the split, update the size to be within
-      // the new block.
-      if (Sym->getOffset() + Sym->getSize() > SplitIndex)
-        Sym->setSize(SplitIndex - Sym->getOffset());
-      Sym->setBlock(NewBlock);
-      BlockSymbols.pop_back();
+
+    auto TransferSymbol = [](Symbol &Sym, Block &B) {
+      Sym.setOffset(Sym.getAddress() - B.getAddress());
+      if (Sym.getSize() > B.getSize())
+        Sym.setSize(B.getSize() - Sym.getOffset());
+      Sym.setBlock(B);
+    };
+
+    // Transfer symbols to all blocks except the last one.
+    for (size_t I = 0; I != Blocks.size() - 1; ++I) {
+      if ((*Cache)->empty())
+        break;
+      while (!(*Cache)->empty() &&
+             (*Cache)->back()->getAddress() < Blocks[I + 1]->getAddress()) {
+        TransferSymbol(*(*Cache)->back(), *Blocks[I]);
+        (*Cache)->pop_back();
+      }
+    }
+    // Transfer symbols to the last block, checking that all are in-range.
+    while (!(*Cache)->empty()) {
+      auto &Sym = *(*Cache)->back();
+      (*Cache)->pop_back();
+      assert(Sym.getAddress() >= Blocks.back()->getAddress() &&
+             "Symbol address preceeds block");
+      assert(Sym.getAddress() <= Blocks.back()->getRange().End &&
+             "Symbol address starts past end of block");
+      TransferSymbol(Sym, *Blocks.back());
+    }
+  }
+
+  // Transfer edges.
+  auto &Edges = Blocks[0]->Edges;
+  llvm::sort(Edges, [](const Edge &LHS, const Edge &RHS) {
+    return LHS.getOffset() < RHS.getOffset();
+  });
+
+  for (size_t I = Blocks.size() - 1; I != 0; --I) {
+
+    // If all edges have been transferred then bail out.
+    if (Edges.empty())
+      break;
+
+    Edge::OffsetT Delta = Blocks[I]->getAddress() - Blocks[0]->getAddress();
+
+    // If no edges to move for this block then move to the next one.
+    if (Edges.back().getOffset() < Delta)
+      continue;
+
+    size_t EI = Edges.size() - 1;
+    while (EI != 0 && Edges[EI - 1].getOffset() >= Delta)
+      --EI;
+
+    for (size_t J = EI; J != Edges.size(); ++J) {
+      Blocks[I]->Edges.push_back(std::move(Edges[J]));
+      Blocks[I]->Edges.back().setOffset(Blocks[I]->Edges.back().getOffset() -
+                                        Delta);
     }
 
-    // Update offsets for all remaining symbols in B.
-    for (auto *Sym : BlockSymbols)
-      Sym->setOffset(Sym->getOffset() - SplitIndex);
+    while (Edges.size() > EI)
+      Edges.pop_back();
   }
 
-  return NewBlock;
+  return Blocks;
 }
 
 void LinkGraph::dump(raw_ostream &OS) {