From 3e78f9852fce9b5de7cef2d6c67fb8fb6b85f395 Mon Sep 17 00:00:00 2001
From: rzvxa <3788964+rzvxa@users.noreply.github.com>
Date: Thu, 20 Jun 2024 01:27:03 +0000
Subject: [PATCH] feat(cfg): add depth first search with hash sets. (#3771)

petgraph allocates a chunk of memory for all of the graph for its dfs which is really costly for small subgraph sreachs.
This PR adds a new `set_depth_first_search` function which uses a `FxHashSet` instead.
---
 crates/oxc_cfg/src/lib.rs   |   7 +--
 crates/oxc_cfg/src/visit.rs | 103 +++++++++++++++++++++++++++++++++++-
 2 files changed, 106 insertions(+), 4 deletions(-)

diff --git a/crates/oxc_cfg/src/lib.rs b/crates/oxc_cfg/src/lib.rs
index e25536f3c2813..af393df6c6e92 100644
--- a/crates/oxc_cfg/src/lib.rs
+++ b/crates/oxc_cfg/src/lib.rs
@@ -6,7 +6,7 @@ use itertools::Itertools;
 use oxc_syntax::node::AstNodeId;
 use petgraph::{
     stable_graph::NodeIndex,
-    visit::{depth_first_search, Control, DfsEvent, EdgeRef},
+    visit::{Control, DfsEvent, EdgeRef},
     Direction, Graph,
 };
 
@@ -21,6 +21,7 @@ pub mod graph {
 
 pub use builder::{ControlFlowGraphBuilder, CtxCursor, CtxFlags};
 pub use dot::DisplayDot;
+use visit::set_depth_first_search;
 
 pub type BasicBlockId = NodeIndex;
 
@@ -155,7 +156,7 @@ impl ControlFlowGraph {
             return true;
         }
         let graph = &self.graph;
-        depth_first_search(&self.graph, Some(from), |event| match event {
+        set_depth_first_search(&self.graph, Some(from), |event| match event {
             DfsEvent::TreeEdge(a, b) => {
                 let filter_result = filter(a);
                 if !matches!(filter_result, Control::Continue) {
@@ -246,7 +247,7 @@ impl ControlFlowGraph {
     }
 
     pub fn is_cyclic(&self, node: BasicBlockId) -> bool {
-        depth_first_search(&self.graph, Some(node), |event| match event {
+        set_depth_first_search(&self.graph, Some(node), |event| match event {
             DfsEvent::BackEdge(_, id) if id == node => Err(()),
             _ => Ok(()),
         })
diff --git a/crates/oxc_cfg/src/visit.rs b/crates/oxc_cfg/src/visit.rs
index 5398643a498a3..72f3d651ab001 100644
--- a/crates/oxc_cfg/src/visit.rs
+++ b/crates/oxc_cfg/src/visit.rs
@@ -1,4 +1,9 @@
-use petgraph::{visit::EdgeRef, Direction, Graph};
+use std::hash::Hash;
+
+use petgraph::{
+    visit::{ControlFlow, DfsEvent, EdgeRef, IntoNeighbors, Time, VisitMap, Visitable},
+    Direction, Graph,
+};
 use rustc_hash::FxHashSet;
 
 use crate::BasicBlockId;
@@ -56,3 +61,99 @@ where
 
     final_states
 }
+
+/// Copied from petgraph's `dfsvisit`.
+/// Return if the expression is a break value, execute the provided statement
+/// if it is a prune value.
+macro_rules! try_control {
+    ($e:expr, $p:stmt) => {
+        try_control!($e, $p, ());
+    };
+    ($e:expr, $p:stmt, $q:stmt) => {
+        match $e {
+            x =>
+            {
+                #[allow(clippy::redundant_else)]
+                if x.should_break() {
+                    return x;
+                } else if x.should_prune() {
+                    $p
+                } else {
+                    $q
+                }
+            }
+        }
+    };
+}
+
+/// Similar to `depth_first_search` but uses a `HashSet` underneath. Ideal for small subgraphs.
+pub fn set_depth_first_search<G, I, F, C, N>(graph: G, starts: I, mut visitor: F) -> C
+where
+    N: Copy + PartialEq + Eq + Hash,
+    G: IntoNeighbors + Visitable<NodeId = N>,
+    I: IntoIterator<Item = G::NodeId>,
+    F: FnMut(DfsEvent<G::NodeId>) -> C,
+    C: ControlFlow,
+{
+    let time = &mut Time(0);
+    let discovered = &mut FxHashSet::<G::NodeId>::default();
+    let finished = &mut FxHashSet::<G::NodeId>::default();
+
+    for start in starts {
+        try_control!(
+            dfs_visitor(graph, start, &mut visitor, discovered, finished, time),
+            unreachable!()
+        );
+    }
+    C::continuing()
+}
+
+fn dfs_visitor<G, M, F, C>(
+    graph: G,
+    u: G::NodeId,
+    visitor: &mut F,
+    discovered: &mut M,
+    finished: &mut M,
+    time: &mut Time,
+) -> C
+where
+    G: IntoNeighbors + Visitable,
+    M: VisitMap<G::NodeId>,
+    F: FnMut(DfsEvent<G::NodeId>) -> C,
+    C: ControlFlow,
+{
+    if !discovered.visit(u) {
+        return C::continuing();
+    }
+
+    try_control!(
+        visitor(DfsEvent::Discover(u, time_post_inc(time))),
+        {},
+        for v in graph.neighbors(u) {
+            if !discovered.is_visited(&v) {
+                try_control!(visitor(DfsEvent::TreeEdge(u, v)), continue);
+                try_control!(
+                    dfs_visitor(graph, v, visitor, discovered, finished, time),
+                    unreachable!()
+                );
+            } else if !finished.is_visited(&v) {
+                try_control!(visitor(DfsEvent::BackEdge(u, v)), continue);
+            } else {
+                try_control!(visitor(DfsEvent::CrossForwardEdge(u, v)), continue);
+            }
+        }
+    );
+    let first_finish = finished.visit(u);
+    debug_assert!(first_finish);
+    try_control!(
+        visitor(DfsEvent::Finish(u, time_post_inc(time))),
+        panic!("Pruning on the `DfsEvent::Finish` is not supported!")
+    );
+    C::continuing()
+}
+
+fn time_post_inc(x: &mut Time) -> Time {
+    let v = *x;
+    x.0 += 1;
+    v
+}