[JITLink] Move block ownership from LinkGraph to Section.

This enables easy iteration over blocks in a specific section.

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

@@ -562,6 +562,16 @@ public:
   /// Returns the ordinal for this section.
   SectionOrdinal getOrdinal() const { return SecOrdinal; }
 
+  /// Returns an iterator over the blocks defined in this section.
+  iterator_range<block_iterator> blocks() {
+    return make_range(Blocks.begin(), Blocks.end());
+  }
+
+  /// Returns an iterator over the blocks defined in this section.
+  iterator_range<const_block_iterator> blocks() const {
+    return make_range(Blocks.begin(), Blocks.end());
+  }
+
   /// Returns an iterator over the symbols defined in this section.
   iterator_range<symbol_iterator> symbols() {
     return make_range(Symbols.begin(), Symbols.end());
@@ -592,10 +602,21 @@ private:
     Symbols.erase(&Sym);
   }
 
+  void addBlock(Block &B) {
+    assert(!Blocks.count(&B) && "Block is already in this section");
+    Blocks.insert(&B);
+  }
+
+  void removeBlock(Block &B) {
+    assert(Blocks.count(&B) && "Block is not in this section");
+    Blocks.erase(&B);
+  }
+
   StringRef Name;
   sys::Memory::ProtectionFlags Prot;
   SectionOrdinal SecOrdinal = 0;
   BlockOrdinal NextBlockOrdinal = 0;
+  BlockSet Blocks;
   SymbolSet Symbols;
 };
 
@@ -663,12 +684,11 @@ private:
   template <typename... ArgTs> Block &createBlock(ArgTs &&... Args) {
     Block *B = reinterpret_cast<Block *>(Allocator.Allocate<Block>());
     new (B) Block(std::forward<ArgTs>(Args)...);
-    Blocks.insert(B);
+    B->getSection().addBlock(*B);
     return *B;
   }
 
   void destroyBlock(Block &B) {
-    Blocks.erase(&B);
     B.~Block();
     Allocator.Deallocate(&B);
   }
@@ -678,69 +698,102 @@ private:
     Allocator.Deallocate(&S);
   }
 
+  static iterator_range<Section::block_iterator> getSectionBlocks(Section &S) {
+    return S.blocks();
+  }
+
+  static iterator_range<Section::const_block_iterator>
+  getSectionConstBlocks(Section &S) {
+    return S.blocks();
+  }
+
+  static iterator_range<Section::symbol_iterator>
+  getSectionSymbols(Section &S) {
+    return S.symbols();
+  }
+
+  static iterator_range<Section::const_symbol_iterator>
+  getSectionConstSymbols(Section &S) {
+    return S.symbols();
+  }
+
 public:
   using external_symbol_iterator = ExternalSymbolSet::iterator;
 
-  using block_iterator = BlockSet::iterator;
-
   using section_iterator = pointee_iterator<SectionList::iterator>;
   using const_section_iterator = pointee_iterator<SectionList::const_iterator>;
 
-  template <typename SectionItrT, typename SymbolItrT, typename T>
+  template <typename OuterItrT, typename InnerItrT, typename T,
-  class defined_symbol_iterator_impl
+            iterator_range<InnerItrT> getInnerRange(
+                typename OuterItrT::reference)>
+  class nested_collection_iterator
       : public iterator_facade_base<
-            defined_symbol_iterator_impl<SectionItrT, SymbolItrT, T>,
+            nested_collection_iterator<OuterItrT, InnerItrT, T, getInnerRange>,
             std::forward_iterator_tag, T> {
   public:
-    defined_symbol_iterator_impl() = default;
+    nested_collection_iterator() = default;
 
-    defined_symbol_iterator_impl(SectionItrT SecI, SectionItrT SecE)
+    nested_collection_iterator(OuterItrT OuterI, OuterItrT OuterE)
-        : SecI(SecI), SecE(SecE),
+        : OuterI(OuterI), OuterE(OuterE),
-          SymI(SecI != SecE ? SecI->symbols().begin() : SymbolItrT()) {
+          InnerI(getInnerBegin(OuterI, OuterE)) {
-      moveToNextSymbolOrEnd();
+      moveToNonEmptyInnerOrEnd();
     }
 
-    bool operator==(const defined_symbol_iterator_impl &RHS) const {
+    bool operator==(const nested_collection_iterator &RHS) const {
-      return (SecI == RHS.SecI) && (SymI == RHS.SymI);
+      return (OuterI == RHS.OuterI) && (InnerI == RHS.InnerI);
     }
 
     T operator*() const {
-      assert(SymI != SecI->symbols().end() && "Dereferencing end?");
+      assert(InnerI != getInnerRange(*OuterI).end() && "Dereferencing end?");
-      return *SymI;
+      return *InnerI;
     }
 
-    defined_symbol_iterator_impl operator++() {
+    nested_collection_iterator operator++() {
-      ++SymI;
+      ++InnerI;
-      moveToNextSymbolOrEnd();
+      moveToNonEmptyInnerOrEnd();
       return *this;
     }
 
   private:
-    void moveToNextSymbolOrEnd() {
+    static InnerItrT getInnerBegin(OuterItrT OuterI, OuterItrT OuterE) {
-      while (SecI != SecE && SymI == SecI->symbols().end()) {
+      return OuterI != OuterE ? getInnerRange(*OuterI).begin() : InnerItrT();
-        ++SecI;
+    }
-        SymI = SecI == SecE ? SymbolItrT() : SecI->symbols().begin();
+
+    void moveToNonEmptyInnerOrEnd() {
+      while (OuterI != OuterE && InnerI == getInnerRange(*OuterI).end()) {
+        ++OuterI;
+        InnerI = getInnerBegin(OuterI, OuterE);
       }
     }
 
-    SectionItrT SecI, SecE;
+    OuterItrT OuterI, OuterE;
-    SymbolItrT SymI;
+    InnerItrT InnerI;
   };
 
   using defined_symbol_iterator =
-      defined_symbol_iterator_impl<const_section_iterator,
+      nested_collection_iterator<const_section_iterator,
-                                   Section::symbol_iterator, Symbol *>;
+                                 Section::symbol_iterator, Symbol *,
+                                 getSectionSymbols>;
 
-  using const_defined_symbol_iterator = defined_symbol_iterator_impl<
+  using const_defined_symbol_iterator =
-      const_section_iterator, Section::const_symbol_iterator, const Symbol *>;
+      nested_collection_iterator<const_section_iterator,
+                                 Section::const_symbol_iterator, const Symbol *,
+                                 getSectionConstSymbols>;
+
+  using block_iterator = nested_collection_iterator<const_section_iterator,
+                                                    Section::block_iterator,
+                                                    Block *, getSectionBlocks>;
+
+  using const_block_iterator =
+      nested_collection_iterator<const_section_iterator,
+                                 Section::const_block_iterator, const Block *,
+                                 getSectionConstBlocks>;
 
   LinkGraph(std::string Name, unsigned PointerSize,
             support::endianness Endianness)
       : Name(std::move(Name)), PointerSize(PointerSize),
         Endianness(Endianness) {}
 
-  ~LinkGraph();
-
   /// Returns the name of this graph (usually the name of the original
   /// underlying MemoryBuffer).
   const std::string &getName() { return Name; }
@@ -864,6 +917,16 @@ public:
     return nullptr;
   }
 
+  iterator_range<block_iterator> blocks() {
+    return make_range(block_iterator(Sections.begin(), Sections.end()),
+                      block_iterator(Sections.end(), Sections.end()));
+  }
+
+  iterator_range<const_block_iterator> blocks() const {
+    return make_range(const_block_iterator(Sections.begin(), Sections.end()),
+                      const_block_iterator(Sections.end(), Sections.end()));
+  }
+
   iterator_range<external_symbol_iterator> external_symbols() {
     return make_range(ExternalSymbols.begin(), ExternalSymbols.end());
   }
@@ -883,10 +946,6 @@ public:
         const_defined_symbol_iterator(Sections.end(), Sections.end()));
   }
 
-  iterator_range<block_iterator> blocks() {
-    return make_range(Blocks.begin(), Blocks.end());
-  }
-
   /// Turn a defined symbol into an external one.
   void makeExternal(Symbol &Sym) {
     if (Sym.getAddressable().isAbsolute()) {
@@ -934,7 +993,7 @@ public:
 
   /// Remove a block.
   void removeBlock(Block &B) {
-    Blocks.erase(&B);
+    B.getSection().removeBlock(B);
     destroyBlock(B);
   }
 
@@ -955,7 +1014,6 @@ private:
   std::string Name;
   unsigned PointerSize;
   support::endianness Endianness;
-  BlockSet Blocks;
   SectionList Sections;
   ExternalSymbolSet ExternalSymbols;
   ExternalSymbolSet AbsoluteSymbols;

llvm/lib/ExecutionEngine/JITLink/JITLink.cpp

@@ -145,10 +145,6 @@ void printEdge(raw_ostream &OS, const Block &B, const Edge &E,
 Section::~Section() {
   for (auto *Sym : Symbols)
     Sym->~Symbol();
-}
-
-LinkGraph::~LinkGraph() {
-  // Destroy blocks.
   for (auto *B : Blocks)
     B->~Block();
 }

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(
+  auto &Sec = G.createSection("__data", RWFlags);
-      "test", sys::Memory::ProtectionFlags(sys::Memory::MF_READ |
-                                           sys::Memory::MF_WRITE));
 
   // Create the block to split.
   auto &B1 = G.createContentBlock(Sec, BlockContent, 0x1000, 8, 0);