[llvm] [ADT] SCC iterator for general graph (PR #84268)

Yida Zhang via llvm-commits llvm-commits at lists.llvm.org
Fri Mar 8 16:36:22 PST 2024


https://github.com/zyd2001 updated https://github.com/llvm/llvm-project/pull/84268

>From 11489913facb5c8362c16bf78766ce33db65fcfa Mon Sep 17 00:00:00 2001
From: Yida Zhang <zyd200101 at outlook.com>
Date: Wed, 6 Mar 2024 20:32:15 -0500
Subject: [PATCH 1/3] SCC iterator for general graph

---
 llvm/include/llvm/ADT/SCCIterator.h | 135 ++++++++++++++++++++++++++++
 1 file changed, 135 insertions(+)

diff --git a/llvm/include/llvm/ADT/SCCIterator.h b/llvm/include/llvm/ADT/SCCIterator.h
index e743ae7c11edbc..0c25323f32f2b3 100644
--- a/llvm/include/llvm/ADT/SCCIterator.h
+++ b/llvm/include/llvm/ADT/SCCIterator.h
@@ -377,6 +377,141 @@ scc_member_iterator<GraphT, GT>::scc_member_iterator(
   assert(InputNodes.size() == Nodes.size() && "missing nodes in MST");
   std::reverse(Nodes.begin(), Nodes.end());
 }
+
+template <class GraphT, class GT = GraphTraits<GraphT>>
+class graph_scc_iterator : public iterator_facade_base<
+                            graph_scc_iterator<GraphT, GT>, std::forward_iterator_tag,
+                            const std::vector<typename GT::NodeRef>, ptrdiff_t> {
+  using NodeRef = typename GT::NodeRef;
+  using NodesIter = typename GT::nodes_iterator;
+  using ChildIter = typename GT::ChildIteratorType;
+  using reference = typename graph_scc_iterator::reference;
+
+  struct NodeEntry {
+    int index = 0;
+    int lowlink;
+    bool onStack = false;
+  };
+
+  struct StackEntry {
+    NodeRef n;
+    ChildIter currentChild;
+    NodeRef lastChild;
+    bool visited;
+  };
+
+  std::vector<NodeRef> currentSCC;
+  std::stack<StackEntry> recursionStack;
+  std::stack<NodeRef> S;
+  DenseMap<NodeRef, NodeEntry> nodeInfo;
+  int index = 1;
+
+  // current node in the outer loop for all nodes
+  NodesIter currentNode;
+  NodesIter nodeEndIter;
+
+  graph_scc_iterator(const GraphT &g, bool e = false)
+      : currentNode(GT::nodes_begin(&g)), nodeEndIter(GT::nodes_end(&g)) {
+    if (e) {
+      currentNode = nodeEndIter;
+      return;
+    }
+    computeNext();
+  }
+
+  void computeNext() {
+    currentSCC.clear();
+    if (recursionStack.empty()) {
+      // not in DFS process
+      for (; currentNode != nodeEndIter; currentNode++) {
+        NodeRef v = *currentNode;
+        if (nodeInfo[v].index == 0) {
+          recursionStack.emplace(StackEntry{v, GT::child_begin(v), nullptr, false});
+          currentNode++;
+          break;
+        }
+      }
+    }
+
+    while (!recursionStack.empty()) {
+      StackEntry &stackEntry = recursionStack.top();
+      NodeRef v = stackEntry.n;
+      ChildIter &child = stackEntry.currentChild;
+      NodeEntry &vEntry = nodeInfo[v];
+      if (!stackEntry.visited) {
+        // first time
+        vEntry.index = index;
+        vEntry.lowlink = index;
+        index++;
+        S.push(v);
+        vEntry.onStack = true;
+        stackEntry.visited = true;
+      } else {
+        assert(nodeInfo.count(stackEntry.lastChild));
+        vEntry.lowlink = std::min(vEntry.lowlink, nodeInfo[stackEntry.lastChild].lowlink);
+      }
+
+      bool inserted = false;
+      for (; child != GT::child_end(v); child++) {
+        NodeRef w = *child;
+        NodeEntry &wEntry = nodeInfo[w];
+        if (wEntry.index == 0) {
+          inserted = true;
+          recursionStack.emplace(StackEntry{w, GT::child_begin(w), nullptr, false});
+          stackEntry.lastChild = w;
+          child++;
+          break;
+        } else if (wEntry.onStack) {
+          vEntry.lowlink = std::min(vEntry.lowlink, wEntry.index);
+        }
+      }
+
+      if (inserted)
+        continue;
+
+      recursionStack.pop();
+      if (vEntry.lowlink == vEntry.index) {
+        NodeRef w;
+        do {
+          w = S.top();
+          S.pop();
+          nodeInfo[w].onStack = false;
+          currentSCC.push_back(w);
+        } while (w != v);
+        return;
+      }
+    }
+  }
+
+public:
+  static graph_scc_iterator begin(const GraphT &g) {
+    return graph_scc_iterator(g);
+  }
+  static graph_scc_iterator end(const GraphT &g) {
+    return graph_scc_iterator(g, true);
+  }
+
+  bool operator!=(const graph_scc_iterator &x) const {
+    return currentNode != x.currentNode || currentSCC != x.currentSCC;
+  }
+
+  graph_scc_iterator &operator++() {
+    computeNext();
+    return *this;
+  }
+
+  reference operator*() const { return currentSCC; }
+};
+/// Construct the begin iterator for a deduced graph type T.
+template <class T> graph_scc_iterator<T> graph_scc_begin(const T &G) {
+  return graph_scc_iterator<T>::begin(G);
+}
+
+/// Construct the end iterator for a deduced graph type T.
+template <class T> graph_scc_iterator<T> graph_scc_end(const T &G) {
+  return graph_scc_iterator<T>::end(G);
+}
+
 } // end namespace llvm
 
 #endif // LLVM_ADT_SCCITERATOR_H

>From 601ad3f8ccbb61fe5f930ddcf584a51fad08be3f Mon Sep 17 00:00:00 2001
From: Yida Zhang <zyd200101 at outlook.com>
Date: Wed, 6 Mar 2024 22:28:02 -0500
Subject: [PATCH 2/3] add header and format

---
 llvm/include/llvm/ADT/SCCIterator.h | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/ADT/SCCIterator.h b/llvm/include/llvm/ADT/SCCIterator.h
index 0c25323f32f2b3..57722789492726 100644
--- a/llvm/include/llvm/ADT/SCCIterator.h
+++ b/llvm/include/llvm/ADT/SCCIterator.h
@@ -31,6 +31,7 @@
 #include <iterator>
 #include <queue>
 #include <set>
+#include <stack>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -379,9 +380,10 @@ scc_member_iterator<GraphT, GT>::scc_member_iterator(
 }
 
 template <class GraphT, class GT = GraphTraits<GraphT>>
-class graph_scc_iterator : public iterator_facade_base<
-                            graph_scc_iterator<GraphT, GT>, std::forward_iterator_tag,
-                            const std::vector<typename GT::NodeRef>, ptrdiff_t> {
+class graph_scc_iterator
+    : public iterator_facade_base<
+          graph_scc_iterator<GraphT, GT>, std::forward_iterator_tag,
+          const std::vector<typename GT::NodeRef>, ptrdiff_t> {
   using NodeRef = typename GT::NodeRef;
   using NodesIter = typename GT::nodes_iterator;
   using ChildIter = typename GT::ChildIteratorType;
@@ -426,7 +428,8 @@ class graph_scc_iterator : public iterator_facade_base<
       for (; currentNode != nodeEndIter; currentNode++) {
         NodeRef v = *currentNode;
         if (nodeInfo[v].index == 0) {
-          recursionStack.emplace(StackEntry{v, GT::child_begin(v), nullptr, false});
+          recursionStack.emplace(
+              StackEntry{v, GT::child_begin(v), nullptr, false});
           currentNode++;
           break;
         }
@@ -448,7 +451,8 @@ class graph_scc_iterator : public iterator_facade_base<
         stackEntry.visited = true;
       } else {
         assert(nodeInfo.count(stackEntry.lastChild));
-        vEntry.lowlink = std::min(vEntry.lowlink, nodeInfo[stackEntry.lastChild].lowlink);
+        vEntry.lowlink =
+            std::min(vEntry.lowlink, nodeInfo[stackEntry.lastChild].lowlink);
       }
 
       bool inserted = false;
@@ -457,7 +461,8 @@ class graph_scc_iterator : public iterator_facade_base<
         NodeEntry &wEntry = nodeInfo[w];
         if (wEntry.index == 0) {
           inserted = true;
-          recursionStack.emplace(StackEntry{w, GT::child_begin(w), nullptr, false});
+          recursionStack.emplace(
+              StackEntry{w, GT::child_begin(w), nullptr, false});
           stackEntry.lastChild = w;
           child++;
           break;

>From bb8ae393c0606032bc2c04568e2dc7faa1b7e1d6 Mon Sep 17 00:00:00 2001
From: Yida Zhang <zyd200101 at outlook.com>
Date: Fri, 8 Mar 2024 19:36:09 -0500
Subject: [PATCH 3/3] Use non-member operator

---
 llvm/include/llvm/ADT/SCCIterator.h | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/ADT/SCCIterator.h b/llvm/include/llvm/ADT/SCCIterator.h
index 57722789492726..8d3307b872eb16 100644
--- a/llvm/include/llvm/ADT/SCCIterator.h
+++ b/llvm/include/llvm/ADT/SCCIterator.h
@@ -413,7 +413,7 @@ class graph_scc_iterator
   NodesIter nodeEndIter;
 
   graph_scc_iterator(const GraphT &g, bool e = false)
-      : currentNode(GT::nodes_begin(&g)), nodeEndIter(GT::nodes_end(&g)) {
+      : currentNode(GT::nodes_begin(g)), nodeEndIter(GT::nodes_end(g)) {
     if (e) {
       currentNode = nodeEndIter;
       return;
@@ -496,9 +496,9 @@ class graph_scc_iterator
     return graph_scc_iterator(g, true);
   }
 
-  bool operator!=(const graph_scc_iterator &x) const {
-    return currentNode != x.currentNode || currentSCC != x.currentSCC;
-  }
+  template <class T>
+  friend bool operator!=(const graph_scc_iterator<T> &x,
+                         const graph_scc_iterator<T> &y);
 
   graph_scc_iterator &operator++() {
     computeNext();
@@ -516,6 +516,11 @@ template <class T> graph_scc_iterator<T> graph_scc_begin(const T &G) {
 template <class T> graph_scc_iterator<T> graph_scc_end(const T &G) {
   return graph_scc_iterator<T>::end(G);
 }
+template <class T>
+bool operator!=(const graph_scc_iterator<T> &x,
+                const graph_scc_iterator<T> &y) {
+  return x.currentNode != y.currentNode || x.currentSCC != y.currentSCC;
+}
 
 } // end namespace llvm
 



More information about the llvm-commits mailing list