Expose recursive defs as a book method and rewrite find_cycles

src/fun/transform/expand_generated.rs

@@ -27,59 +27,61 @@ type DepGraph = HashMap<Name, HashSet<Name>>;
 type Cycles = Vec<Vec<Name>>;
 type RecursiveDefs = BTreeSet<Name>;
 
-pub fn find_recursive_defs(book: &Book) -> RecursiveDefs {
-  let mut cycle_map = BTreeSet::new();
-  let deps = book_def_deps(book);
-  let cycles = cycles(&deps);
+impl Book {
+  pub fn recursive_defs(&self) -> RecursiveDefs {
+    let mut cycle_map = BTreeSet::new();
+    let deps = book_def_deps(self);
+    let cycles = cycles(&deps);
 
-  for cycle in cycles {
-    for name in cycle {
-      cycle_map.insert(name);
+    for cycle in cycles {
+      for name in cycle {
+        cycle_map.insert(name);
+      }
     }
-  }
 
-  cycle_map
+    cycle_map
+  }
 }
 
 /// Find all cycles in the dependency graph.
 fn cycles(deps: &DepGraph) -> Cycles {
   let mut cycles = vec![];
-  let mut stack = vec![];
   let mut visited = HashSet::new();
   for nam in deps.keys() {
     if !visited.contains(nam) {
-      find_cycles(deps, nam, &mut visited, &mut stack, &mut cycles);
+      find_cycles(deps, nam, &mut visited, &mut cycles);
     }
   }
   cycles
 }
 
-fn find_cycles(
-  deps: &DepGraph,
-  nam: &Name,
-  visited: &mut HashSet<Name>,
-  stack: &mut Vec<Name>,
-  cycles: &mut Cycles,
-) {
+fn find_cycles(deps: &DepGraph, nam: &Name, visited: &mut HashSet<Name>, cycles: &mut Cycles) {
   maybe_grow(|| {
-    // Check if the current ref is already in the stack, which indicates a cycle.
-    if let Some(cycle_start) = stack.iter().position(|n| n == nam) {
-      // If found, add the cycle to the cycles vector.
-      cycles.push(stack[cycle_start..].to_vec());
-      return;
-    }
-    // If the ref has not been visited yet, mark it as visited.
-    if visited.insert(nam.clone()) {
-      // Add the current ref to the stack to keep track of the path.
-      stack.push(nam.clone());
-      // Get the dependencies of the current ref.
-      if let Some(dependencies) = deps.get(nam) {
-        // Search for cycles from each dependency.
-        for dep in dependencies {
-          find_cycles(deps, dep, visited, stack, cycles);
+    let mut to_search = vec![(nam.clone(), false)];
+    while let Some((current, checked)) = to_search.pop() {
+      if checked {
+        to_search.pop();
+      } else {
+        // Check if the current ref is already in the stack, which indicates a cycle.
+        if let Some(cycle_start) = to_search.iter().position(|(n, _)| n == &current) {
+          // If found, add the cycle to the cycles vector.
+          cycles.push(to_search[cycle_start..].iter().map(|(n, _)| n.clone()).collect::<Vec<_>>());
+          continue;
+        }
+        // If the ref has not been visited yet, mark it as visited.
+        if visited.insert(current.clone()) {
+          // Get the dependencies of the current ref.
+          let deps = deps.get(&current);
+          // Add the current ref to the stack to keep track of the path.
+          to_search.push((current, true));
+          if let Some(deps) = deps {
+            // Search for cycles from each dependency.
+            for dep in deps {
+              to_search.push((dep.clone(), false));
+            }
+          }
         }
       }
-      stack.pop();
     }
   })
 }

src/lib.rs

@@ -11,7 +11,6 @@ use crate::{
   },
 };
 use diagnostics::{Diagnostics, DiagnosticsConfig, ERR_INDENT_SIZE};
-use fun::transform::expand_generated::find_recursive_defs;
 use net::hvm_to_net::hvm_to_net;
 
 pub mod diagnostics;
@@ -192,7 +191,7 @@ pub fn readback_hvm_net(
   let mut diags = Diagnostics::default();
   let net = hvm_to_net(net);
   let mut term = net_to_term(&net, book, labels, linear, &mut diags);
-  let recursive_cycles = find_recursive_defs(book);
+  let recursive_cycles = book.recursive_defs();
   term.expand_generated(book, &recursive_cycles);
   term.resugar_strings(adt_encoding);
   term.resugar_lists(adt_encoding);