[llvm] 521c996 - [JITLink] Move Symbol to new block before updating size.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2024-11-18T10:38:50+11:00
New Revision: 521c99627690e09ba25383c83232f94ff457f00c

URL: https://github.com/llvm/llvm-project/commit/521c99627690e09ba25383c83232f94ff457f00c
DIFF: https://github.com/llvm/llvm-project/commit/521c99627690e09ba25383c83232f94ff457f00c.diff

LOG: [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).

Added: 
    

Modified: 
    llvm/lib/ExecutionEngine/JITLink/JITLink.cpp
    llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLink.cpp
index ef382c3ce695a9..ba30a76934e117 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 32d917d75d5ca4..936559f072e052 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) {