diff --git a/hugr-passes/src/monomorphize.rs b/hugr-passes/src/monomorphize.rs index d7a9bb7ad..5d5eafca6 100644 --- a/hugr-passes/src/monomorphize.rs +++ b/hugr-passes/src/monomorphize.rs @@ -10,7 +10,7 @@ use hugr_core::{ Node, }; -use hugr_core::hugr::{hugrmut::HugrMut, internal::HugrMutInternals, Hugr, HugrView, OpType}; +use hugr_core::hugr::{hugrmut::HugrMut, Hugr, HugrView, OpType}; use itertools::Itertools as _; /// Replaces calls to polymorphic functions with calls to new monomorphic @@ -31,23 +31,25 @@ use itertools::Itertools as _; pub fn monomorphize(mut h: Hugr) -> Hugr { let validate = |h: &Hugr| h.validate().unwrap_or_else(|e| panic!("{e}")); - // We clone the extension registry because we will need a reference to - // create our mutable substitutions. This is cannot cause a problem because - // we will not be adding any new types or extension ops to the HUGR. #[cfg(debug_assertions)] validate(&h); + monomorphize_ref(&mut h); + + #[cfg(debug_assertions)] + validate(&h); + h +} + +fn monomorphize_ref(h: &mut impl HugrMut) { let root = h.root(); // If the root is a polymorphic function, then there are no external calls, so nothing to do if !is_polymorphic_funcdefn(h.get_optype(root)) { - mono_scan(&mut h, root, None, &mut HashMap::new()); + mono_scan(h, root, None, &mut HashMap::new()); if !h.get_optype(root).is_module() { - return remove_polyfuncs(h); + remove_polyfuncs_ref(h); } } - #[cfg(debug_assertions)] - validate(&h); - h } /// Removes any polymorphic [FuncDefn]s from the Hugr. Note that if these have @@ -72,6 +74,21 @@ pub fn remove_polyfuncs(mut h: Hugr) -> Hugr { h } +fn remove_polyfuncs_ref(h: &mut impl HugrMut) { + let mut pfs_to_delete = Vec::new(); + let mut to_scan = Vec::from_iter(h.children(h.root())); + while let Some(n) = to_scan.pop() { + if is_polymorphic_funcdefn(h.get_optype(n)) { + pfs_to_delete.push(n) + } else { + to_scan.extend(h.children(n)); + } + } + for n in pfs_to_delete { + h.remove_subtree(n); + } +} + fn is_polymorphic(fd: &FuncDefn) -> bool { !fd.signature.params().is_empty() } @@ -93,7 +110,7 @@ type Instantiations = HashMap, Node>>; /// Optionally copies the subtree into a new location whilst applying a substitution. /// The subtree should be monomorphic after the substitution (if provided) has been applied. fn mono_scan( - h: &mut Hugr, + h: &mut impl HugrMut, parent: Node, mut subst_into: Option<&mut Instantiating>, cache: &mut Instantiations, @@ -161,7 +178,7 @@ fn mono_scan( } fn instantiate( - h: &mut Hugr, + h: &mut impl HugrMut, poly_func: Node, type_args: Vec, mono_sig: Signature,