[JITLink] Allow moving a Symbol to a Block in a different Section

[benlangmuir]

Summary

When moving a Symbol between Blocks that are in different Sections, update the symbol tables for each Section. Otherwise symbol.getBlock().getSection() will not match the contents of Section::symbols(), which asserts during linking.

[lhames]

LGTM. Thanks very much Ben!

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

public:
   /// If ExplicitNewSize is given as None then the size of the symbol will be
   /// checked and auto-truncated to at most the size of the remainder (from the
   /// given offset) of the size of the new block.
   ///
   /// All other symbol attributes are unchanged.
   void transferDefinedSymbol(Symbol &Sym, Block &DestBlock,
                              JITTargetAddress NewOffset,
                              Optional<JITTargetAddress> ExplicitNewSize) {
+    auto &OldSection = Sym.getBlock().getSection();
     Sym.setBlock(DestBlock);
     Sym.setOffset(NewOffset);
     if (ExplicitNewSize)
       Sym.setSize(*ExplicitNewSize);
     else {
       JITTargetAddress RemainingBlockSize = DestBlock.getSize() - NewOffset;
       if (Sym.getSize() > RemainingBlockSize)
         Sym.setSize(RemainingBlockSize);
     }
+    if (&DestBlock.getSection() != &OldSection) {
+      OldSection.removeSymbol(Sym);
+      DestBlock.getSection().addSymbol(Sym);
+    }
   }
 
   /// 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

llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

TEST(LinkGraphTest, TransferDefinedSymbol) {
   // Transfer with non-zero offset, implicit truncation.
   G.transferDefinedSymbol(S1, B3, 16, None);
 
   EXPECT_EQ(&S1.getBlock(), &B3) << "Block was not updated";
   EXPECT_EQ(S1.getOffset(), 16U) << "Offset was not updated";
   EXPECT_EQ(S1.getSize(), 16U) << "Size was not updated";
 }
 
+TEST(LinkGraphTest, TransferDefinedSymbolAcrossSections) {
+  // Check that we can transfer a defined symbol from an existing block in 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 blocks in each section.
+  auto &B1 = G.createContentBlock(Sec1, BlockContent, 0x1000, 8, 0);
+  auto &B2 = G.createContentBlock(Sec2, BlockContent, 0x2000, 8, 0);
+
+  // Add a symbol to section 1.
+  auto &S1 = G.addDefinedSymbol(B1, 0, "S1", B1.getSize(), Linkage::Strong,
+                                Scope::Default, false, false);
+
+  // Transfer with zero offset, explicit size to section 2.
+  G.transferDefinedSymbol(S1, B2, 0, 64);
+
+  EXPECT_EQ(&S1.getBlock(), &B2) << "Block was not updated";
+  EXPECT_EQ(S1.getOffset(), 0U) << "Unexpected offset";
+  EXPECT_EQ(S1.getSize(), 64U) << "Size was not updated";
+
+  EXPECT_EQ(Sec1.symbols_size(), 0u) << "Symbol was not removed from Sec1";
+  EXPECT_EQ(Sec2.symbols_size(), 1u) << "Symbol was not added to Sec2";
+  if (Sec2.symbols_size() == 1)
+    EXPECT_EQ(*Sec2.symbols().begin(), &S1) << "Unexpected symbol";
+}
+
 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);