[llvm] 29e6332 - [JITLink] Add Block::edges_at(Edge::OffsetT): iterate over edges at offset.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2025-01-16T13:18:15+11:00
New Revision: 29e63328a79af9501bf8d6b7e5a07303427ded73

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

LOG: [JITLink] Add Block::edges_at(Edge::OffsetT): iterate over edges at offset.

Block::edges_at is a convenience method for iterating over edges at a given
offset within a jitlink::Block.

This method will be used in an upcoming patch for compact unwind info support.

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 67bcb007873121..bd82fadea10270 100644
--- a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
+++ b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
@@ -336,6 +336,18 @@ class Block : public Addressable {
     return make_range(Edges.begin(), Edges.end());
   }
 
+  /// Returns an iterator over all edges at the given offset within the block.
+  auto edges_at(Edge::OffsetT O) {
+    return make_filter_range(edges(),
+                             [O](const Edge &E) { return E.getOffset() == O; });
+  }
+
+  /// Returns an iterator over all edges at the given offset within the block.
+  auto edges_at(Edge::OffsetT O) const {
+    return make_filter_range(edges(),
+                             [O](const Edge &E) { return E.getOffset() == O; });
+  }
+
   /// Return the size of the edges list.
   size_t edges_size() const { return Edges.size(); }
 

diff  --git a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
index ff6cf49bb97587..fb60acddf7821e 100644
--- a/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
+++ b/llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp
@@ -127,6 +127,36 @@ TEST(LinkGraphTest, BlockAndSymbolIteration) {
   EXPECT_TRUE(llvm::count(G.defined_symbols(), &S4));
 }
 
+TEST(LinkGraphTest, EdgeIteration) {
+  // Check that we can iterate over blocks within Sections and across sections.
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+  auto &Sec1 =
+      G.createSection("__data.1", orc::MemProt::Read | orc::MemProt::Write);
+  auto &B =
+      G.createContentBlock(Sec1, BlockContent, orc::ExecutorAddr(0x1000), 8, 0);
+  auto &S = G.addExternalSymbol("S1", 0, false);
+
+  constexpr size_t NumEdges = 6;
+  Edge::OffsetT Offsets[NumEdges] = {0, 1, 2, 2, 3, 7};
+
+  for (auto O : Offsets)
+    B.addEdge(Edge::KeepAlive, O, S, 0);
+
+  EXPECT_EQ(llvm::range_size(B.edges()), NumEdges);
+  EXPECT_EQ(llvm::range_size(B.edges_at(0)), 1U);
+  EXPECT_EQ(llvm::range_size(B.edges_at(2)), 2U);
+  EXPECT_EQ(llvm::range_size(B.edges_at(4)), 0U);
+
+  {
+    // Check that offsets and iteration order are as expected.
+    size_t Idx = 0;
+    for (auto &E : B.edges())
+      EXPECT_EQ(E.getOffset(), Offsets[Idx++]);
+  }
+}
+
 TEST(LinkGraphTest, ContentAccessAndUpdate) {
   // Check that we can make a defined symbol external.
   LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),