[JITLink] Add support for moving blocks and symbols between sections.

LinkGraph::transferBlock can be used to move a block and all associated symbols
from one section to another.

LinkGraph::mergeSections moves all blocks and sections from a source section to
a destination section.

include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -154,7 +154,7 @@ private:
   /// Create a zero-fill defined addressable.
   Block(Section &Parent, JITTargetAddress Size, JITTargetAddress Address,
         uint64_t Alignment, uint64_t AlignmentOffset)
-      : Addressable(Address, true), Parent(Parent), Size(Size) {
+      : Addressable(Address, true), Parent(&Parent), Size(Size) {
     assert(isPowerOf2_64(Alignment) && "Alignment must be power of 2");
     assert(AlignmentOffset < Alignment &&
            "Alignment offset cannot exceed alignment");
@@ -170,7 +170,7 @@ private:
   /// mutations are required.
   Block(Section &Parent, ArrayRef<char> Content, JITTargetAddress Address,
         uint64_t Alignment, uint64_t AlignmentOffset)
-      : Addressable(Address, true), Parent(Parent), Data(Content.data()),
+      : Addressable(Address, true), Parent(&Parent), Data(Content.data()),
         Size(Content.size()) {
     assert(isPowerOf2_64(Alignment) && "Alignment must be power of 2");
     assert(AlignmentOffset < Alignment &&
@@ -189,7 +189,7 @@ private:
   /// allocator.
   Block(Section &Parent, MutableArrayRef<char> Content,
         JITTargetAddress Address, uint64_t Alignment, uint64_t AlignmentOffset)
-      : Addressable(Address, true), Parent(Parent), Data(Content.data()),
+      : Addressable(Address, true), Parent(&Parent), Data(Content.data()),
         Size(Content.size()) {
     assert(isPowerOf2_64(Alignment) && "Alignment must be power of 2");
     assert(AlignmentOffset < Alignment &&
@@ -212,7 +212,7 @@ public:
   Block &operator=(Block &&) = delete;
 
   /// Return the parent section for this block.
-  Section &getSection() const { return Parent; }
+  Section &getSection() const { return *Parent; }
 
   /// Returns true if this is a zero-fill block.
   ///
@@ -331,7 +331,9 @@ public:
 private:
   static constexpr uint64_t MaxAlignmentOffset = (1ULL << 56) - 1;
 
-  Section &Parent;
+  void setSection(Section &Parent) { this->Parent = &Parent; }
+
+  Section *Parent;
   const char *Data = nullptr;
   size_t Size = 0;
   std::vector<Edge> Edges;
@@ -684,6 +686,8 @@ public:
     return make_range(Blocks.begin(), Blocks.end());
   }
 
+  BlockSet::size_type blocks_size() const { return Blocks.size(); }
+
   /// Returns an iterator over the symbols defined in this section.
   iterator_range<symbol_iterator> symbols() {
     return make_range(Symbols.begin(), Symbols.end());
@@ -695,7 +699,7 @@ public:
   }
 
   /// Return the number of symbols in this section.
-  SymbolSet::size_type symbols_size() { return Symbols.size(); }
+  SymbolSet::size_type symbols_size() const { return Symbols.size(); }
 
 private:
   void addSymbol(Symbol &Sym) {
@@ -718,6 +722,17 @@ private:
     Blocks.erase(&B);
   }
 
+  void transferContentTo(Section &DstSection) {
+    if (&DstSection == this)
+      return;
+    for (auto *S : Symbols)
+      DstSection.addSymbol(*S);
+    for (auto *B : Blocks)
+      DstSection.addBlock(*B);
+    Symbols.clear();
+    Blocks.clear();
+  }
+
   StringRef Name;
   sys::Memory::ProtectionFlags Prot;
   SectionOrdinal SecOrdinal = 0;
@@ -1102,6 +1117,8 @@ public:
                       section_iterator(Sections.end()));
   }
 
+  SectionList::size_type sections_size() const { return Sections.size(); }
+
   /// Returns the section with the given name if it exists, otherwise returns
   /// null.
   Section *findSectionByName(StringRef Name) {
@@ -1231,6 +1248,43 @@ public:
     }
   }
 
+  /// Transfers the given Block and all Symbols pointing to it to the given
+  /// Section.
+  ///
+  /// No attempt is made to check compatibility of the source and destination
+  /// sections. Blocks may be moved between sections with incompatible
+  /// permissions (e.g. from data to text). The client is responsible for
+  /// ensuring that this is safe.
+  void transferBlock(Block &B, Section &NewSection) {
+    auto &OldSection = B.getSection();
+    if (&OldSection == &NewSection)
+      return;
+    SmallVector<Symbol *> AttachedSymbols;
+    for (auto *S : OldSection.symbols())
+      if (&S->getBlock() == &B)
+        AttachedSymbols.push_back(S);
+    for (auto *S : AttachedSymbols) {
+      OldSection.removeSymbol(*S);
+      NewSection.addSymbol(*S);
+    }
+    OldSection.removeBlock(B);
+    NewSection.addBlock(B);
+  }
+
+  /// Move all blocks and symbols from the source section to the destination
+  /// section.
+  ///
+  /// If PreserveSrcSection is true (or SrcSection and DstSection are the same)
+  /// then SrcSection is preserved, otherwise it is removed (the default).
+  void mergeSections(Section &DstSection, Section &SrcSection,
+                     bool PreserveSrcSection = false) {
+    if (&DstSection == &SrcSection)
+      return;
+    SrcSection.transferContentTo(DstSection);
+    if (!PreserveSrcSection)
+      removeSection(SrcSection);
+  }
+
   /// Removes an external symbol. Also removes the underlying Addressable.
   void removeExternalSymbol(Symbol &Sym) {
     assert(!Sym.isDefined() && !Sym.isAbsolute() &&
@@ -1269,7 +1323,8 @@ public:
     destroySymbol(Sym);
   }
 
-  /// Remove a block.
+  /// Remove a block. The block reference is defunct after calling this
+  /// function and should no longer be used.
   void removeBlock(Block &B) {
     assert(llvm::none_of(B.getSection().symbols(),
                          [&](const Symbol *Sym) {
@@ -1280,6 +1335,16 @@ public:
     destroyBlock(B);
   }
 
+  /// Remove a section. The section reference is defunct after calling this
+  /// function and should no longer be used.
+  void removeSection(Section &Sec) {
+    auto I = llvm::find_if(Sections, [&Sec](const std::unique_ptr<Section> &S) {
+      return S.get() == &Sec;
+    });
+    assert(I != Sections.end() && "Section does not appear in this graph");
+    Sections.erase(I);
+  }
+
   /// Dump the graph.
   void dump(raw_ostream &OS);
 

lib/ExecutionEngine/JITLink/JITLink.cpp

@@ -264,7 +264,10 @@ void LinkGraph::dump(raw_ostream &OS) {
       OS << "  block " << formatv("{0:x16}", B->getAddress())
          << " size = " << formatv("{0:x8}", B->getSize())
          << ", align = " << B->getAlignment()
-         << ", alignment-offset = " << B->getAlignmentOffset() << "\n";
+         << ", alignment-offset = " << B->getAlignmentOffset();
+      if (B->isZeroFill())
+        OS << ", zero-fill";
+      OS << "\n";
 
       auto BlockSymsI = BlockSymbols.find(B);
       if (BlockSymsI != BlockSymbols.end()) {

unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

@@ -323,6 +323,125 @@ TEST(LinkGraphTest, TransferDefinedSymbol) {
   EXPECT_EQ(S1.getSize(), 16U) << "Size was not updated";
 }
 
+TEST(LinkGraphTest, TransferBlock) {
+  // Check that we can transfer a block (and all associated symbols) from one
+  // section to another.
+  LinkGraph G("foo", Triple("x86_64-apple-darwin"), 8, support::little,
+              getGenericEdgeKindName);
+  auto &Sec1 = G.createSection("__data.1", RWFlags);
+  auto &Sec2 = G.createSection("__data.2", RWFlags);
+
+  // Create an initial block.
+  auto &B1 = G.createContentBlock(Sec1, BlockContent, 0x1000, 8, 0);
+  auto &B2 = G.createContentBlock(Sec1, BlockContent, 0x2000, 8, 0);
+
+  // Add some symbols on B1...
+  G.addDefinedSymbol(B1, 0, "S1", B1.getSize(), Linkage::Strong, Scope::Default,
+                     false, false);
+  G.addDefinedSymbol(B1, 1, "S2", B1.getSize() - 1, Linkage::Strong,
+                     Scope::Default, false, false);
+
+  // ... and on B2.
+  G.addDefinedSymbol(B2, 0, "S3", B2.getSize(), Linkage::Strong, Scope::Default,
+                     false, false);
+  G.addDefinedSymbol(B2, 1, "S4", B2.getSize() - 1, Linkage::Strong,
+                     Scope::Default, false, false);
+
+  EXPECT_EQ(Sec1.blocks_size(), 2U) << "Expected two blocks in Sec1 initially";
+  EXPECT_EQ(Sec1.symbols_size(), 4U)
+      << "Expected four symbols in Sec1 initially";
+  EXPECT_EQ(Sec2.blocks_size(), 0U) << "Expected zero blocks in Sec2 initially";
+  EXPECT_EQ(Sec2.symbols_size(), 0U)
+      << "Expected zero symbols in Sec2 initially";
+
+  // Transfer with zero offset, explicit size.
+  G.transferBlock(B1, Sec2);
+
+  EXPECT_EQ(Sec1.blocks_size(), 1U)
+      << "Expected one blocks in Sec1 after transfer";
+  EXPECT_EQ(Sec1.symbols_size(), 2U)
+      << "Expected two symbols in Sec1 after transfer";
+  EXPECT_EQ(Sec2.blocks_size(), 1U)
+      << "Expected one blocks in Sec2 after transfer";
+  EXPECT_EQ(Sec2.symbols_size(), 2U)
+      << "Expected two symbols in Sec2 after transfer";
+}
+
+TEST(LinkGraphTest, MergeSections) {
+  // Check that we can transfer a block (and all associated symbols) from one
+  // section to another.
+  LinkGraph G("foo", Triple("x86_64-apple-darwin"), 8, support::little,
+              getGenericEdgeKindName);
+  auto &Sec1 = G.createSection("__data.1", RWFlags);
+  auto &Sec2 = G.createSection("__data.2", RWFlags);
+  auto &Sec3 = G.createSection("__data.3", RWFlags);
+
+  // Create an initial block.
+  auto &B1 = G.createContentBlock(Sec1, BlockContent, 0x1000, 8, 0);
+  auto &B2 = G.createContentBlock(Sec2, BlockContent, 0x2000, 8, 0);
+  auto &B3 = G.createContentBlock(Sec3, BlockContent, 0x3000, 8, 0);
+
+  // Add a symbols for each block.
+  G.addDefinedSymbol(B1, 0, "S1", B1.getSize(), Linkage::Strong, Scope::Default,
+                     false, false);
+  G.addDefinedSymbol(B2, 0, "S2", B2.getSize(), Linkage::Strong, Scope::Default,
+                     false, false);
+  G.addDefinedSymbol(B3, 0, "S3", B2.getSize(), Linkage::Strong, Scope::Default,
+                     false, false);
+
+  EXPECT_EQ(G.sections_size(), 3U) << "Expected three sections initially";
+  EXPECT_EQ(Sec1.blocks_size(), 1U) << "Expected one block in Sec1 initially";
+  EXPECT_EQ(Sec1.symbols_size(), 1U) << "Expected one symbol in Sec1 initially";
+  EXPECT_EQ(Sec2.blocks_size(), 1U) << "Expected one block in Sec2 initially";
+  EXPECT_EQ(Sec2.symbols_size(), 1U) << "Expected one symbol in Sec2 initially";
+  EXPECT_EQ(Sec3.blocks_size(), 1U) << "Expected one block in Sec3 initially";
+  EXPECT_EQ(Sec3.symbols_size(), 1U) << "Expected one symbol in Sec3 initially";
+
+  // Check that self-merge is a no-op.
+  G.mergeSections(Sec1, Sec1);
+
+  EXPECT_EQ(G.sections_size(), 3U)
+      << "Expected three sections after first merge";
+  EXPECT_EQ(Sec1.blocks_size(), 1U)
+      << "Expected one block in Sec1 after first merge";
+  EXPECT_EQ(Sec1.symbols_size(), 1U)
+      << "Expected one symbol in Sec1 after first merge";
+  EXPECT_EQ(Sec2.blocks_size(), 1U)
+      << "Expected one block in Sec2 after first merge";
+  EXPECT_EQ(Sec2.symbols_size(), 1U)
+      << "Expected one symbol in Sec2 after first merge";
+  EXPECT_EQ(Sec3.blocks_size(), 1U)
+      << "Expected one block in Sec3 after first merge";
+  EXPECT_EQ(Sec3.symbols_size(), 1U)
+      << "Expected one symbol in Sec3 after first merge";
+
+  // Merge Sec2 into Sec1, removing Sec2.
+  G.mergeSections(Sec1, Sec2);
+
+  EXPECT_EQ(G.sections_size(), 2U)
+      << "Expected two sections after section merge";
+  EXPECT_EQ(Sec1.blocks_size(), 2U)
+      << "Expected two blocks in Sec1 after section merge";
+  EXPECT_EQ(Sec1.symbols_size(), 2U)
+      << "Expected two symbols in Sec1 after section merge";
+  EXPECT_EQ(Sec3.blocks_size(), 1U)
+      << "Expected one block in Sec3 after section merge";
+  EXPECT_EQ(Sec3.symbols_size(), 1U)
+      << "Expected one symbol in Sec3 after section merge";
+
+  G.mergeSections(Sec1, Sec3, true);
+
+  EXPECT_EQ(G.sections_size(), 2U) << "Expected two sections after third merge";
+  EXPECT_EQ(Sec1.blocks_size(), 3U)
+      << "Expected three blocks in Sec1 after third merge";
+  EXPECT_EQ(Sec1.symbols_size(), 3U)
+      << "Expected three symbols in Sec1 after third merge";
+  EXPECT_EQ(Sec3.blocks_size(), 0U)
+      << "Expected one block in Sec3 after third merge";
+  EXPECT_EQ(Sec3.symbols_size(), 0U)
+      << "Expected one symbol in Sec3 after third merge";
+}
+
 TEST(LinkGraphTest, SplitBlock) {
   // Check that the LinkGraph::splitBlock test works as expected.
   LinkGraph G("foo", Triple("x86_64-apple-darwin"), 8, support::little,