[JITLink] Move block ownership from LinkGraph to Section.

This enables easy iteration over blocks in a specific section.

1 files changed, 57 insertions, 3 deletions

diff --git a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
index cb86d09f023..c1e1d9e740d 100644
--- a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
+++ b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
@@ -15,6 +15,9 @@
 using namespace llvm;
 using namespace llvm::jitlink;
 
+static auto RWFlags =
+    sys::Memory::ProtectionFlags(sys::Memory::MF_READ | sys::Memory::MF_WRITE);
+
 TEST(LinkGraphTest, Construction) {
   // Check that LinkGraph construction works as expected.
   LinkGraph G("foo", 8, support::little);
@@ -27,6 +30,59 @@ TEST(LinkGraphTest, Construction) {
   EXPECT_TRUE(llvm::empty(G.blocks()));
 }
 
+TEST(LinkGraphTest, BlockAndSymbolIteration) {
+  // Check that we can iterate over blocks within Sections and across sections.
+
+  const char BlockContentBytes[] = {0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+                                    0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
+                                    0x1C, 0x1D, 0x1E, 0x1F, 0x00};
+  StringRef BlockContent(BlockContentBytes);
+
+  LinkGraph G("foo", 8, support::little);
+  auto &Sec1 = G.createSection("__data.1", RWFlags);
+  auto &B1 = G.createContentBlock(Sec1, BlockContent, 0x1000, 8, 0);
+  auto &B2 = G.createContentBlock(Sec1, BlockContent, 0x2000, 8, 0);
+  auto &S1 = G.addDefinedSymbol(B1, 0, "S1", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+  auto &S2 = G.addDefinedSymbol(B2, 4, "S2", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+
+  auto &Sec2 = G.createSection("__data.1", RWFlags);
+  auto &B3 = G.createContentBlock(Sec2, BlockContent, 0x3000, 8, 0);
+  auto &B4 = G.createContentBlock(Sec2, BlockContent, 0x4000, 8, 0);
+  auto &S3 = G.addDefinedSymbol(B3, 0, "S3", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+  auto &S4 = G.addDefinedSymbol(B4, 4, "S4", 4, Linkage::Strong, Scope::Default,
+                                false, false);
+
+  // Check that iteration of blocks within a section behaves as expected.
+  EXPECT_EQ(std::distance(Sec1.blocks().begin(), Sec1.blocks().end()), 2U);
+  EXPECT_TRUE(llvm::count(Sec1.blocks(), &B1));
+  EXPECT_TRUE(llvm::count(Sec1.blocks(), &B2));
+
+  // Check that iteration of symbols within a section behaves as expected.
+  EXPECT_EQ(std::distance(Sec1.symbols().begin(), Sec1.symbols().end()), 2U);
+  EXPECT_TRUE(llvm::count(Sec1.symbols(), &S1));
+  EXPECT_TRUE(llvm::count(Sec1.symbols(), &S2));
+
+  // Check that iteration of blocks across sections behaves as expected.
+  EXPECT_EQ(std::distance(G.blocks().begin(), G.blocks().end()), 4U);
+  EXPECT_TRUE(llvm::count(G.blocks(), &B1));
+  EXPECT_TRUE(llvm::count(G.blocks(), &B2));
+  EXPECT_TRUE(llvm::count(G.blocks(), &B3));
+  EXPECT_TRUE(llvm::count(G.blocks(), &B4));
+
+  // Check that iteration of defined symbols across sections behaves as
+  // expected.
+  EXPECT_EQ(
+      std::distance(G.defined_symbols().begin(), G.defined_symbols().end()),
+      4U);
+  EXPECT_TRUE(llvm::count(G.defined_symbols(), &S1));
+  EXPECT_TRUE(llvm::count(G.defined_symbols(), &S2));
+  EXPECT_TRUE(llvm::count(G.defined_symbols(), &S3));
+  EXPECT_TRUE(llvm::count(G.defined_symbols(), &S4));
+}
+
 TEST(LinkGraphTest, SplitBlock) {
   // Check that the LinkGraph::splitBlock test works as expected.
 
@@ -36,9 +92,7 @@ TEST(LinkGraphTest, SplitBlock) {
   StringRef BlockContent(BlockContentBytes);
 
   LinkGraph G("foo", 8, support::little);
-  auto &Sec = G.createSection(
-      "test", sys::Memory::ProtectionFlags(sys::Memory::MF_READ |
-                                           sys::Memory::MF_WRITE));
+  auto &Sec = G.createSection("__data", RWFlags);
 
   // Create the block to split.
   auto &B1 = G.createContentBlock(Sec, BlockContent, 0x1000, 8, 0);