[llvm] b24ec07 - [JITLink] Update Section pointers on Blocks when merging Sections.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2021-11-10T22:53:03-08:00
New Revision: b24ec07f43a5793b3b1912b29b17967ef642d30e

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

LOG: [JITLink] Update Section pointers on Blocks when merging Sections.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
    llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
index a9e99547f005f..4a5d18437e39e 100644
--- a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
+++ b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
@@ -1299,6 +1299,8 @@ class LinkGraph {
                      bool PreserveSrcSection = false) {
     if (&DstSection == &SrcSection)
       return;
+    for (auto *B : SrcSection.blocks())
+      B->setSection(DstSection);
     SrcSection.transferContentTo(DstSection);
     if (!PreserveSrcSection)
       removeSection(SrcSection);

diff  --git a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
index c8958e4d9f3b8..4851f642e85a5 100644
--- a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
+++ b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
@@ -416,6 +416,8 @@ TEST(LinkGraphTest, MergeSections) {
   G.addDefinedSymbol(B3, 0, "S3", B2.getSize(), Linkage::Strong, Scope::Default,
                      false, false);
 
+  EXPECT_EQ(&B1.getSection(), &Sec1);
+  EXPECT_EQ(&B2.getSection(), &Sec2);
   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";
@@ -427,6 +429,8 @@ TEST(LinkGraphTest, MergeSections) {
   // Check that self-merge is a no-op.
   G.mergeSections(Sec1, Sec1);
 
+  EXPECT_EQ(&B1.getSection(), &Sec1)
+      << "Expected B1.getSection() to remain unchanged";
   EXPECT_EQ(G.sections_size(), 3U)
       << "Expected three sections after first merge";
   EXPECT_EQ(Sec1.blocks_size(), 1U)
@@ -445,6 +449,8 @@ TEST(LinkGraphTest, MergeSections) {
   // Merge Sec2 into Sec1, removing Sec2.
   G.mergeSections(Sec1, Sec2);
 
+  EXPECT_EQ(&B2.getSection(), &Sec1)
+      << "Expected B2.getSection() to have been changed to &Sec1";
   EXPECT_EQ(G.sections_size(), 2U)
       << "Expected two sections after section merge";
   EXPECT_EQ(Sec1.blocks_size(), 2U)