diff --git a/src/fun/transform/expand_generated.rs b/src/fun/transform/expand_generated.rs index f801d325..910b42a7 100644 --- a/src/fun/transform/expand_generated.rs +++ b/src/fun/transform/expand_generated.rs @@ -27,59 +27,61 @@ type DepGraph = HashMap>; type Cycles = Vec>; type RecursiveDefs = BTreeSet; -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, - stack: &mut Vec, - cycles: &mut Cycles, -) { +fn find_cycles(deps: &DepGraph, nam: &Name, visited: &mut HashSet, 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 == ¤t) { + // If found, add the cycle to the cycles vector. + cycles.push(to_search[cycle_start..].iter().map(|(n, _)| n.clone()).collect::>()); + 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(¤t); + // 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(); } }) } diff --git a/src/lib.rs b/src/lib.rs index f3d63924..fc58055b 100644 --- a/src/lib.rs +++ b/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);