[JITLink] Move Symbol to new block before updating size.

Symbol::setSize asserts that the new size does not overflow the containing
block, so we need to point the Symbol at the correct Block before updating its
size (otherwise we may get a spurious overflow assertion).

2 files changed, 38 insertions, 29 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp
index ef382c3..ba30a76 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp
@@ -204,9 +204,9 @@ std::vector<Block *> LinkGraph::splitBlockImpl(std::vector<Block *> Blocks,
 
     auto TransferSymbol = [](Symbol &Sym, Block &B) {
       Sym.setOffset(Sym.getAddress() - B.getAddress());
+      Sym.setBlock(B);
       if (Sym.getSize() > B.getSize())
         Sym.setSize(B.getSize() - Sym.getOffset());
-      Sym.setBlock(B);
     };
 
     // Transfer symbols to all blocks except the last one.
diff --git a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
index 32d917d..936559f 100644
--- a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
+++ b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
@@ -622,9 +622,12 @@ TEST(LinkGraphTest, SplitBlock) {
                                 false, false);
   auto &S4 = G.addDefinedSymbol(B1, 12, "S4", 4, Linkage::Strong,
                                 Scope::Default, false, false);
-  // Add a symbol that extends beyond the split.
+  // Add some symbols that extend beyond splits, one in the first block and one
+  // in a subsequent block.
   auto &S5 = G.addDefinedSymbol(B1, 0, "S5", 16, Linkage::Strong,
                                 Scope::Default, false, false);
+  auto &S6 = G.addDefinedSymbol(B1, 6, "S6", 10, Linkage::Strong,
+                                Scope::Default, false, false);
 
   // Add an extra block, EB, and target symbols, and use these to add edges
   // from B1 to EB.
@@ -646,54 +649,56 @@ TEST(LinkGraphTest, SplitBlock) {
   B1.addEdge(Edge::FirstRelocation, 12, ES4, 0);
 
   // Split B1.
-  auto Blocks = G.splitBlock(B1, ArrayRef<int>({8}));
+  auto Blocks = G.splitBlock(B1, ArrayRef<int>({4, 12}));
 
-  EXPECT_EQ(Blocks.size(), 2U);
+  EXPECT_EQ(Blocks.size(), 3U);
   EXPECT_EQ(Blocks[0], &B1);
   auto &B2 = *Blocks[1];
+  auto &B3 = *Blocks[2];
 
   // Check that the block addresses and content matches what we would expect.
   EXPECT_EQ(B1.getAddress(), B1Addr);
-  EXPECT_EQ(B1.getContent(), BlockContent.slice(0, 8));
-  EXPECT_EQ(B1.edges_size(), 2U);
+  EXPECT_EQ(B1.getContent(), BlockContent.slice(0, 4));
+  EXPECT_EQ(B1.edges_size(), 1U);
 
-  EXPECT_EQ(B2.getAddress(), B1Addr + 8);
-  EXPECT_EQ(B2.getContent(), BlockContent.slice(8));
+  EXPECT_EQ(B2.getAddress(), B1Addr + 4);
+  EXPECT_EQ(B2.getContent(), BlockContent.slice(4, 8));
   EXPECT_EQ(B2.edges_size(), 2U);
 
+  EXPECT_EQ(B3.getAddress(), B1Addr + 12);
+  EXPECT_EQ(B3.getContent(), BlockContent.slice(12));
+  EXPECT_EQ(B3.edges_size(), 1U);
+
   // Check that symbols in B2 were transferred as expected:
-  // We expect S1 and S2 to have been transferred to B1, and S3 and S4 to have
-  // remained attached to B2. Symbols S3 and S4 should have had their offsets
-  // slid to account for the change in address of B1.
+  // We expect S1 and S5 to have been transferred to B1; S2, S3 and S6 to
+  // B2; and S4 to B3. Symbols should have had their offsets slid to account
+  // for the change of containing block.
   EXPECT_EQ(&S1.getBlock(), &B1);
   EXPECT_EQ(S1.getOffset(), 0U);
 
-  EXPECT_EQ(&S2.getBlock(), &B1);
-  EXPECT_EQ(S2.getOffset(), 4U);
+  EXPECT_EQ(&S2.getBlock(), &B2);
+  EXPECT_EQ(S2.getOffset(), 0U);
 
   EXPECT_EQ(&S3.getBlock(), &B2);
-  EXPECT_EQ(S3.getOffset(), 0U);
+  EXPECT_EQ(S3.getOffset(), 4U);
 
-  EXPECT_EQ(&S4.getBlock(), &B2);
-  EXPECT_EQ(S4.getOffset(), 4U);
+  EXPECT_EQ(&S4.getBlock(), &B3);
+  EXPECT_EQ(S4.getOffset(), 0U);
 
   EXPECT_EQ(&S5.getBlock(), &B1);
   EXPECT_EQ(S5.getOffset(), 0U);
 
+  EXPECT_EQ(&S6.getBlock(), &B2);
+  EXPECT_EQ(S6.getOffset(), 2U);
+
   // Size shrinks to fit.
-  EXPECT_EQ(S5.getSize(), 8U);
-
-  // Check that edges in B2 have been transferred as expected:
-  // Both blocks should now have two edges each at offsets 0 and 4.
-  EXPECT_EQ(llvm::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(S5.getSize(), 4U);
+  EXPECT_EQ(S6.getSize(), 6U);
+
+  // Check that edges in have been transferred as expected:
+  EXPECT_EQ(llvm::size(B1.edges()), 1);
+  if (size(B1.edges()) == 2)
+    EXPECT_EQ(B1.edges().begin()->getOffset(), 0U);
 
   EXPECT_EQ(llvm::size(B2.edges()), 2);
   if (size(B2.edges()) == 2) {
@@ -704,6 +709,10 @@ TEST(LinkGraphTest, SplitBlock) {
     EXPECT_EQ(E1->getOffset(), 0U);
     EXPECT_EQ(E2->getOffset(), 4U);
   }
+
+  EXPECT_EQ(llvm::size(B3.edges()), 1);
+  if (size(B3.edges()) == 2)
+    EXPECT_EQ(B3.edges().begin()->getOffset(), 0U);
 }
 
 TEST(LinkGraphTest, GraphAllocationMethods) {