Made feedback set algorithm not ignore the nodes post cycle. (#5539)

crates/cairo-lang-lowering/src/graph_algorithms/feedback_set.rs

@@ -41,7 +41,7 @@ pub fn priv_function_with_body_feedback_set_of_representative(
     function: ConcreteSCCRepresentative,
 ) -> Maybe<OrderedHashSet<ConcreteFunctionWithBodyId>> {
     Ok(calc_feedback_set(
-        &ConcreteFunctionWithBodyNode {
+        ConcreteFunctionWithBodyNode {
             function_id: function.0,
             db,
             dependency_type: DependencyType::Cost,

crates/cairo-lang-utils/src/graph_algos/feedback_set.rs

@@ -7,6 +7,8 @@
 //! so here we implement some straight-forward algorithm that guarantees to cover all the cycles in
 //! the graph, but doesn't necessarily produce the minimum size of such a set.
 
+use std::collections::VecDeque;
+
 use super::graph_node::GraphNode;
 use super::scc_graph_node::SccGraphNode;
 use super::strongly_connected_components::ComputeScc;
@@ -19,45 +21,55 @@ mod feedback_set_test;
 
 /// Context for the feedback-set algorithm.
 struct FeedbackSetAlgoContext<Node: ComputeScc> {
+    /// Nodes that were discovered as reachable from the current run, but possibly were not yet
+    /// visited.
+    pending: VecDeque<SccGraphNode<Node>>,
     /// The accumulated feedback set so far in the process of the algorithm. In the end of the
     /// algorithm, this is also the result.
-    pub feedback_set: OrderedHashSet<Node::NodeId>,
+    feedback_set: OrderedHashSet<Node::NodeId>,
     /// Nodes that are currently during the recursion call on them. That is - if one of these is
     /// reached, it indicates it's in some cycle that was not "resolved" yet.
-    pub in_flight: UnorderedHashSet<Node::NodeId>,
-}
-impl<Node: ComputeScc> FeedbackSetAlgoContext<Node> {
-    fn new() -> Self {
-        FeedbackSetAlgoContext {
-            feedback_set: OrderedHashSet::default(),
-            in_flight: UnorderedHashSet::default(),
-        }
-    }
+    in_flight: UnorderedHashSet<Node::NodeId>,
+    /// Already visited nodes in the current run.
+    visited: UnorderedHashSet<Node::NodeId>,
 }
 
 /// Calculates the feedback set of an SCC.
 pub fn calc_feedback_set<Node: ComputeScc>(
-    node: &SccGraphNode<Node>,
+    node: SccGraphNode<Node>,
 ) -> OrderedHashSet<Node::NodeId> {
-    let mut ctx = FeedbackSetAlgoContext::<Node>::new();
-    calc_feedback_set_recursive(node, &mut ctx);
+    let mut ctx = FeedbackSetAlgoContext {
+        feedback_set: OrderedHashSet::default(),
+        in_flight: UnorderedHashSet::default(),
+        pending: VecDeque::new(),
+        visited: UnorderedHashSet::default(),
+    };
+    ctx.pending.push_back(node);
+    while let Some(node) = ctx.pending.pop_front() {
+        calc_feedback_set_recursive(node, &mut ctx);
+    }
     ctx.feedback_set
 }
 
 fn calc_feedback_set_recursive<Node: ComputeScc>(
-    node: &SccGraphNode<Node>,
+    node: SccGraphNode<Node>,
     ctx: &mut FeedbackSetAlgoContext<Node>,
 ) {
     let cur_node_id = node.get_id();
+    if ctx.visited.contains(&cur_node_id) {
+        return;
+    }
+    ctx.visited.insert(cur_node_id.clone());
     ctx.in_flight.insert(cur_node_id.clone());
-    for neighbor in node.get_neighbors() {
+    let mut neighbors = node.get_neighbors().into_iter();
+    for neighbor in neighbors.by_ref() {
         let neighbor_id = neighbor.get_id();
         if ctx.feedback_set.contains(&neighbor_id) {
             continue;
         } else if ctx.in_flight.contains(&neighbor_id) {
             ctx.feedback_set.insert(neighbor_id);
         } else {
-            calc_feedback_set_recursive(&neighbor, ctx);
+            calc_feedback_set_recursive(neighbor, ctx);
         }
 
         // `node` might have been added to the fset during this iteration of the loop. If so, no
@@ -66,5 +78,6 @@ fn calc_feedback_set_recursive<Node: ComputeScc>(
             break;
         }
     }
+    ctx.pending.extend(neighbors.filter(|node| !ctx.visited.contains(&node.get_id())));
     ctx.in_flight.remove(&cur_node_id);
 }

crates/cairo-lang-utils/src/graph_algos/feedback_set_test.rs

@@ -41,7 +41,7 @@ fn test_list() {
     let mut graph: Vec<Vec<usize>> = (0..10).map(|id| vec![id + 1]).collect();
     graph.push(vec![]);
 
-    let fset = HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: 0, graph }.into()));
+    let fset = HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: 0, graph }.into()));
     assert!(fset.is_empty());
 }
 
@@ -50,7 +50,7 @@ fn test_cycle() {
     // Each node has only one neighbor. i -> i + 1 for i = 0...8, and 9 -> 0.
     let graph: Vec<Vec<usize>> = (0..10).map(|id| vec![(id + 1) % 10]).collect();
 
-    let fset = HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: 0, graph }.into()));
+    let fset = HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: 0, graph }.into()));
     assert_eq!(fset, HashSet::from([0]));
 }
 
@@ -62,7 +62,7 @@ fn test_root_points_to_cycle() {
     graph.push(/* 10: */ vec![0]);
 
     // Note 10 is used as a root.
-    let fset = HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: 0, graph }.into()));
+    let fset = HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: 0, graph }.into()));
     assert_eq!(fset, HashSet::from([0]));
 }
 
@@ -77,7 +77,7 @@ fn test_connected_cycles() {
     graph[4].push(5);
 
     // Make sure the cycle that's not in the SCC of the root is not covered.
-    let fset = HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: 0, graph }.into()));
+    let fset = HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: 0, graph }.into()));
     assert_eq!(fset, HashSet::from([0]));
 }
 
@@ -86,7 +86,7 @@ fn test_mesh() {
     // Each node has edges to all other nodes.
     let graph = (0..5).map(|i| (0..5).filter(|j| *j != i).collect::<Vec<usize>>()).collect();
 
-    let fset = HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: 0, graph }.into()));
+    let fset = HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: 0, graph }.into()));
     assert_eq!(fset, HashSet::from_iter(0..4));
 }
 
@@ -100,13 +100,13 @@ fn test_tangent_cycles(root: usize, expected_fset: HashSet<usize>) {
     let graph: Vec<Vec<usize>> = chain!(
         (0..3).map(|id| vec![id + 1]),
         // 3:
-        vec![vec![0, 4]].into_iter(),
+        [vec![0, 4]],
         (0..3).map(|id| vec![3 + (id + 2) % 4])
     )
     .collect();
 
     let fset =
-        HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: root, graph }.into()));
+        HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: root, graph }.into()));
     assert_eq!(fset, expected_fset);
 }
 
@@ -116,19 +116,37 @@ fn test_tangent_cycles(root: usize, expected_fset: HashSet<usize>) {
 #[test_case(2, HashSet::from([2, 3]); "root_2")]
 #[test_case(3, HashSet::from([3]); "root_3")]
 fn test_multiple_cycles(root: usize, expected_fset: HashSet<usize>) {
-    let graph: Vec<Vec<usize>> = chain!(
+    let graph: Vec<Vec<usize>> = vec![
         // 0:
-        vec![vec![1, 2]].into_iter(),
+        vec![1, 2],
         // 1:
-        vec![vec![2, 3]].into_iter(),
+        vec![2, 3],
         // 2:
-        vec![vec![3]].into_iter(),
+        vec![3],
         // 3:
-        vec![vec![0]].into_iter(),
-    )
-    .collect();
+        vec![0],
+    ];
+
+    let fset =
+        HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: root, graph }.into()));
+    assert_eq!(fset, expected_fset);
+}
+
+// Test a graph and continue from self loops.
+#[test_case(0, HashSet::from([1, 2]); "root_0")]
+#[test_case(1, HashSet::from([1, 2]); "root_1")]
+#[test_case(2, HashSet::from([1, 2]); "root_2")]
+fn test_with_self_loops(root: usize, expected_fset: HashSet<usize>) {
+    let graph: Vec<Vec<usize>> = vec![
+        // 0:
+        vec![1, 2],
+        // 1:
+        vec![1, 0],
+        // 2:
+        vec![2, 0],
+    ];
 
     let fset =
-        HashSet::<usize>::from_iter(calc_feedback_set(&IntegerNode { id: root, graph }.into()));
+        HashSet::<usize>::from_iter(calc_feedback_set(IntegerNode { id: root, graph }.into()));
     assert_eq!(fset, expected_fset);
 }