[JITLink] Add a utility for splitting blocks at a given index.

LinkGraph::splitBlock will split a block at a given index, returning a new
block covering the range [ 0, index ) and modifying the original block to
cover the range [ index, original-block-size ). Block addresses, content,
edges and symbols will be updated as necessary. This utility will be used
in upcoming improvements to JITLink's eh-frame support.

include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -72,6 +72,7 @@ public:
       : Target(&Target), Offset(Offset), Addend(Addend), K(K) {}
 
   OffsetT getOffset() const { return Offset; }
+  void setOffset(OffsetT Offset) { this->Offset = Offset; }
   Kind getKind() const { return K; }
   void setKind(Kind K) { this->K = K; }
   bool isRelocation() const { return K >= FirstRelocation; }
@@ -208,15 +209,32 @@ public:
   /// Get the alignment for this content.
   uint64_t getAlignment() const { return 1ull << P2Align; }
 
+  /// Set the alignment for this content.
+  void setAlignment(uint64_t Alignment) {
+    assert(isPowerOf2_64(Alignment) && "Alignment must be a power of two");
+    P2Align = Alignment ? countTrailingZeros(Alignment) : 0;
+  }
+
   /// Get the alignment offset for this content.
   uint64_t getAlignmentOffset() const { return AlignmentOffset; }
 
+  /// Set the alignment offset for this content.
+  void setAlignmentOffset(uint64_t AlignmentOffset) {
+    assert(AlignmentOffset < (1ull << P2Align) &&
+           "Alignment offset can't exceed alignment");
+    this->AlignmentOffset = AlignmentOffset;
+  }
+
   /// Add an edge to this block.
   void addEdge(Edge::Kind K, Edge::OffsetT Offset, Symbol &Target,
                Edge::AddendT Addend) {
     Edges.push_back(Edge(K, Offset, Target, Addend));
   }
 
+  /// Add an edge by copying an existing one. This is typically used when
+  /// moving edges between blocks.
+  void addEdge(const Edge &E) { Edges.push_back(E); }
+
   /// Return the list of edges attached to this content.
   iterator_range<edge_iterator> edges() {
     return make_range(Edges.begin(), Edges.end());
@@ -233,6 +251,10 @@ public:
   /// Returns true if the list of edges is empty.
   bool edges_empty() const { return Edges.empty(); }
 
+  /// Remove the edge pointed to by the given iterator.
+  /// Invalidates all iterators that point to or past the given one.
+  void removeEdge(const_edge_iterator I) { Edges.erase(I); }
+
 private:
   static constexpr uint64_t MaxAlignmentOffset = (1ULL << 57) - 1;
 
@@ -287,6 +309,7 @@ private:
          JITTargetAddress Size, Linkage L, Scope S, bool IsLive,
          bool IsCallable)
       : Name(Name), Base(&Base), Offset(Offset), Size(Size) {
+    assert(Offset <= MaxOffset && "Offset out of range");
     setLinkage(L);
     setScope(S);
     setLive(IsLive);
@@ -484,6 +507,13 @@ private:
     // note: Size and IsCallable fields left unchanged.
   }
 
+  void setBlock(Block &B) { Base = &B; }
+
+  void setOffset(uint64_t NewOffset) {
+    assert(NewOffset <= MaxOffset && "Offset out of range");
+    Offset = NewOffset;
+  }
+
   static constexpr uint64_t MaxOffset = (1ULL << 59) - 1;
 
   // FIXME: A char* or SymbolStringPtr may pack better.
@@ -743,6 +773,29 @@ public:
                        Alignment, AlignmentOffset);
   }
 
+  /// Cache type for the splitBlock function.
+  using SplitBlockCache = 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() ).
+  ///
+  /// 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.
+  ///
+  /// Note: 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);
+
   /// Add an external symbol.
   /// Some formats (e.g. ELF) allow Symbols to have sizes. For Symbols whose
   /// size is not known, you should substitute '0'.

lib/ExecutionEngine/JITLink/JITLink.cpp

@@ -153,6 +153,85 @@ LinkGraph::~LinkGraph() {
     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().substr(0, SplitIndex),
+                B.getAddress(), B.getAlignment(), B.getAlignmentOffset());
+
+  // Modify B to cover [ SplitIndex, B.size() ).
+  B.setAddress(B.getAddress() + SplitIndex);
+  B.setContent(B.getContent().substr(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(), E = B.edges().end(); I != E; ++I) {
+      if (I->getOffset() < SplitIndex) {
+        NewBlock.addEdge(*I);
+        EdgesToRemove.push_back(I);
+      } else
+        I->setOffset(I->getOffset() - SplitIndex);
+    }
+
+    // Remove edges that were transfered to NewBlock from B.
+    while (!EdgesToRemove.empty()) {
+      B.removeEdge(EdgesToRemove.back());
+      EdgesToRemove.pop_back();
+    }
+  }
+
+  // Handle symbol transfer/update.
+  {
+    // Initialize the symbols cache if necessary.
+    SplitBlockCache LocalBlockSymbolsCache;
+    if (!Cache)
+      Cache = &LocalBlockSymbolsCache;
+    if (*Cache == None) {
+      *Cache = SplitBlockCache::value_type();
+      for (auto *Sym : B.getSection().symbols())
+        if (&Sym->getBlock() == &B)
+          (*Cache)->push_back(Sym);
+
+      llvm::sort(**Cache, [](const Symbol *LHS, const Symbol *RHS) {
+        return LHS->getOffset() > RHS->getOffset();
+      });
+    }
+    auto &BlockSymbols = **Cache;
+
+    // Transfer all symbols with offset less than SplitIndex to NewBlock.
+    while (!BlockSymbols.empty() &&
+           BlockSymbols.back()->getOffset() < SplitIndex) {
+      BlockSymbols.back()->setBlock(NewBlock);
+      BlockSymbols.pop_back();
+    }
+
+    // Update offsets for all remaining symbols in B.
+    for (auto *Sym : BlockSymbols)
+      Sym->setOffset(Sym->getOffset() - SplitIndex);
+  }
+
+  return NewBlock;
+}
+
 void LinkGraph::dump(raw_ostream &OS,
                      std::function<StringRef(Edge::Kind)> EdgeKindToName) {
   if (!EdgeKindToName)

unittests/ExecutionEngine/JITLink/CMakeLists.txt

@@ -12,6 +12,7 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest(JITLinkTests
     JITLinkTestCommon.cpp
+    LinkGraphTests.cpp
     MachO_x86_64_Tests.cpp
   )
 

unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

@@ -0,0 +1,120 @@
+//===------ LinkGraphTests.cpp - Unit tests for core JITLink classes ------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/Memory.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+
+TEST(LinkGraphTest, Construction) {
+  // Check that LinkGraph construction works as expected.
+  LinkGraph G("foo", 8, support::little);
+  EXPECT_EQ(G.getName(), "foo");
+  EXPECT_EQ(G.getPointerSize(), 8U);
+  EXPECT_EQ(G.getEndianness(), support::little);
+  EXPECT_TRUE(empty(G.external_symbols()));
+  EXPECT_TRUE(empty(G.absolute_symbols()));
+  EXPECT_TRUE(empty(G.defined_symbols()));
+  EXPECT_TRUE(empty(G.blocks()));
+}
+
+TEST(LinkGraphTest, SplitBlock) {
+  // Check that the LinkGraph::splitBlock test works as expected.
+
+  const char BlockContentBytes[] = {0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+                                    0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
+                                    0x1C, 0x1D, 0x1E, 0x1F, 0x00};
+  StringRef BlockContent(BlockContentBytes);
+
+  LinkGraph G("foo", 8, support::little);
+  auto &Sec = G.createSection(
+      "test", sys::Memory::ProtectionFlags(sys::Memory::MF_READ |
+                                           sys::Memory::MF_WRITE));
+
+  // Create the block to split.
+  auto &B1 = G.createContentBlock(Sec, BlockContent, 0x1000, 8, 0);
+
+  // Add some symbols to the block.
+  auto &S1 = G.addDefinedSymbol(B1, 0, "S1", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+  auto &S2 = G.addDefinedSymbol(B1, 4, "S2", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+  auto &S3 = G.addDefinedSymbol(B1, 8, "S3", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+  auto &S4 = G.addDefinedSymbol(B1, 12, "S4", 4, Linkage::Strong,
+                                Scope::Default, false, false);
+
+  // Add an extra block, EB, and target symbols, and use these to add edges
+  // from B1 to EB.
+  auto &EB = G.createContentBlock(Sec, BlockContent, 0x2000, 8, 0);
+  auto &ES1 = G.addDefinedSymbol(EB, 0, "TS1", 4, Linkage::Strong,
+                                 Scope::Default, false, false);
+  auto &ES2 = G.addDefinedSymbol(EB, 4, "TS2", 4, Linkage::Strong,
+                                 Scope::Default, false, false);
+  auto &ES3 = G.addDefinedSymbol(EB, 8, "TS3", 4, Linkage::Strong,
+                                 Scope::Default, false, false);
+  auto &ES4 = G.addDefinedSymbol(EB, 12, "TS4", 4, Linkage::Strong,
+                                 Scope::Default, false, false);
+
+  // Add edges from B1 to EB.
+  B1.addEdge(Edge::FirstRelocation, 0, ES1, 0);
+  B1.addEdge(Edge::FirstRelocation, 4, ES2, 0);
+  B1.addEdge(Edge::FirstRelocation, 8, ES3, 0);
+  B1.addEdge(Edge::FirstRelocation, 12, ES4, 0);
+
+  // Split B1.
+  auto &B2 = G.splitBlock(B1, 8);
+
+  // Check that the block addresses and content matches what we would expect.
+  EXPECT_EQ(B1.getAddress(), 0x1008U);
+  EXPECT_EQ(B1.getContent(), BlockContent.substr(8));
+
+  EXPECT_EQ(B2.getAddress(), 0x1000U);
+  EXPECT_EQ(B2.getContent(), BlockContent.substr(0, 8));
+
+  // Check that symbols in B1 were transferred as expected:
+  // We expect S1 and S2 to have been transferred to B2, and S3 and S4 to have
+  // remained attached to B1. Symbols S3 and S4 should have had their offsets
+  // slid to account for the change in address of B2.
+  EXPECT_EQ(&S1.getBlock(), &B2);
+  EXPECT_EQ(S1.getOffset(), 0U);
+
+  EXPECT_EQ(&S2.getBlock(), &B2);
+  EXPECT_EQ(S2.getOffset(), 4U);
+
+  EXPECT_EQ(&S3.getBlock(), &B1);
+  EXPECT_EQ(S3.getOffset(), 0U);
+
+  EXPECT_EQ(&S4.getBlock(), &B1);
+  EXPECT_EQ(S4.getOffset(), 4U);
+
+  // Check that edges in B1 have been transferred as expected:
+  // Both blocks should now have two edges each at offsets 0 and 4.
+  EXPECT_EQ(size(B1.edges()), 2);
+  if (size(B1.edges()) == 2) {
+    auto *E1 = &*B1.edges().begin();
+    auto *E2 = &*(B1.edges().begin() + 1);
+    if (E2->getOffset() < E1->getOffset())
+      std::swap(E1, E2);
+    EXPECT_EQ(E1->getOffset(), 0U);
+    EXPECT_EQ(E2->getOffset(), 4U);
+  }
+
+  EXPECT_EQ(size(B2.edges()), 2);
+  if (size(B2.edges()) == 2) {
+    auto *E1 = &*B2.edges().begin();
+    auto *E2 = &*(B2.edges().begin() + 1);
+    if (E2->getOffset() < E1->getOffset())
+      std::swap(E1, E2);
+    EXPECT_EQ(E1->getOffset(), 0U);
+    EXPECT_EQ(E2->getOffset(), 4U);
+  }
+}