Include in DeadCode/DeadFunc elimination

acl-cqc · acl-cqc · commit aaeffa6e125c · 2025-05-12T18:52:33.000+01:00
diff --git a/hugr-passes/src/composable.rs b/hugr-passes/src/composable.rs
@@ -246,15 +246,16 @@ mod test {
             .define_function("id1", Signature::new_endo(usize_t()))
             .unwrap();
         let inps = id1.input_wires();
-        let id1 = id1.finish_with_outputs(inps).unwrap();
+        id1.finish_with_outputs(inps).unwrap();
+
         let id2 = mb
-            .define_function("id2", Signature::new_endo(usize_t()))
+            .define_function_link_name("id2", Signature::new_endo(usize_t()), None)
             .unwrap();
         let inps = id2.input_wires();
         let id2 = id2.finish_with_outputs(inps).unwrap();
         let hugr = mb.finish_hugr().unwrap();
 
-        let dce = DeadCodeElimPass::default().with_entry_points([id1.node()]);
+        let dce = DeadCodeElimPass::default();
         let cfold =
             ConstantFoldPass::default().with_inputs(id2.node(), [(0, ConstUsize::new(2).into())]);
 
diff --git a/hugr-passes/src/dead_code.rs b/hugr-passes/src/dead_code.rs
@@ -20,13 +20,15 @@ pub struct DeadCodeElimPass<H: HugrView> {
     /// Callback identifying nodes that must be preserved even if their
     /// results are not used. Defaults to [PreserveNode::default_for].
     preserve_callback: Arc<PreserveCallback<H>>,
+    include_exports: bool,
 }
 
 impl<H: HugrView + 'static> Default for DeadCodeElimPass<H> {
     fn default() -> Self {
         Self {
             entry_points: Default::default(),
             preserve_callback: Arc::new(PreserveNode::default_for),
+            include_exports: true,
         }
     }
 }
@@ -94,18 +96,33 @@ impl<H: HugrView> DeadCodeElimPass<H> {
     /// Mark some nodes as entry points to the Hugr, i.e. so we cannot eliminate any code
     /// used to evaluate these nodes.
     /// [`HugrView::entrypoint`] is assumed to be an entry point;
-    /// for Module roots the client will want to mark some of the FuncDefn children
-    /// as entry points too.
+    /// if the entrypoint is a Module, then any public
+    /// [FuncDefn](OpType::FuncDefn)s and [Const](OpType::Const)s are also considered entry points
+    /// by default, but these can be removed by [Self::include_module_exports]`(false)`.
     pub fn with_entry_points(mut self, entry_points: impl IntoIterator<Item = H::Node>) -> Self {
         self.entry_points.extend(entry_points);
         self
     }
 
+    /// Sets whether, for Module-rooted Hugrs, the exported [FuncDefn](OpType::FuncDefn)s
+    /// and [Const](OpType::Const)s are included as entry points (they are by default)
+    pub fn include_module_exports(mut self, include: bool) -> Self {
+        self.include_exports = include;
+        self
+    }
+
     fn find_needed_nodes(&self, h: &H) -> HashSet<H::Node> {
         let mut must_preserve = HashMap::new();
         let mut needed = HashSet::new();
         let mut q = VecDeque::from_iter(self.entry_points.iter().cloned());
         q.push_front(h.entrypoint());
+        if self.include_exports && h.entrypoint() == h.module_root() {
+            q.extend(h.children(h.module_root()).filter(|ch| {
+                h.get_optype(*ch)
+                    .as_func_defn()
+                    .is_some_and(|fd| fd.link_name.is_some())
+            }))
+        }
         while let Some(n) = q.pop_front() {
             if !needed.insert(n) {
                 continue;
@@ -120,8 +137,8 @@ impl<H: HugrView> DeadCodeElimPass<H> {
                     | OpType::AliasDecl(_) // and all Aliases (we do not track their uses in types)
                     | OpType::AliasDefn(_)
                     | OpType::Input(_) // Also Dataflow input/output, these are necessary for legality
-                    | OpType::Output(_) // Do not include FuncDecl / FuncDefn / Const unless reachable by static edges
-                                                                // (from Call/LoadConst/LoadFunction):
+                    | OpType::Output(_) // Do not include FuncDecl / Const unless reachable by static edges
+                                                                // (from Call/LoadConst/LoadFunction)
                     )
                 {
                     q.push_back(ch);
diff --git a/hugr-passes/src/dead_funcs.rs b/hugr-passes/src/dead_funcs.rs
@@ -61,10 +61,20 @@ fn reachable_funcs<'a, H: HugrView>(
     }))
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 /// A configuration for the Dead Function Removal pass.
 pub struct RemoveDeadFuncsPass {
     entry_points: Vec<Node>,
+    include_exports: bool,
+}
+
+impl Default for RemoveDeadFuncsPass {
+    fn default() -> Self {
+        Self {
+            entry_points: Default::default(),
+            include_exports: true,
+        }
+    }
 }
 
 impl RemoveDeadFuncsPass {
@@ -80,16 +90,34 @@ impl RemoveDeadFuncsPass {
         self.entry_points.extend(entry_points);
         self
     }
+
+    /// Sets whether the exported [FuncDefn](hugr_core::ops::FuncDefn) children of a
+    /// [Module](hugr_core::ops::Module) are included as entry points (yes by default)
+    pub fn include_module_exports(mut self, include: bool) -> Self {
+        self.include_exports = include;
+        self
+    }
 }
 
 impl<H: HugrMut<Node = Node>> ComposablePass<H> for RemoveDeadFuncsPass {
     type Error = RemoveDeadFuncsError;
     type Result = ();
     fn run(&self, hugr: &mut H) -> Result<(), RemoveDeadFuncsError> {
+        let exports = if hugr.entrypoint() == hugr.module_root() && self.include_exports {
+            hugr.children(hugr.module_root())
+                .filter(|ch| {
+                    hugr.get_optype(*ch)
+                        .as_func_defn()
+                        .is_some_and(|fd| fd.link_name.is_some())
+                })
+                .collect()
+        } else {
+            vec![]
+        };
         let reachable = reachable_funcs(
             &CallGraph::new(hugr),
             hugr,
-            self.entry_points.iter().cloned(),
+            self.entry_points.iter().cloned().chain(exports),
         )?
         .collect::<HashSet<_>>();
         let unreachable = hugr
@@ -108,30 +136,13 @@ impl<H: HugrMut<Node = Node>> ComposablePass<H> for RemoveDeadFuncsPass {
 /// Deletes from the Hugr any functions that are not used by either [Call] or
 /// [LoadFunction] nodes in reachable parts.
 ///
-/// For [Module]-rooted Hugrs, `entry_points` may provide a list of entry points,
-/// which must be children of the root. Note that if `entry_points` is empty, this will
-/// result in all functions in the module being removed.
-///
-/// For non-[Module]-rooted Hugrs, `entry_points` must be empty; the root node is used.
-///
-/// # Errors
-/// * If there are any `entry_points` but the root of the hugr is not a [Module]
-/// * If any node in `entry_points` is
-///     * not a [FuncDefn], or
-///     * not a child of the root
+/// For [Module]-rooted Hugrs, all top-level functions with [FuncDefn::link_name] set,
+/// will be used as entry points.
 ///
-/// [Call]: hugr_core::ops::OpType::Call
-/// [FuncDefn]: hugr_core::ops::OpType::FuncDefn
-/// [LoadFunction]: hugr_core::ops::OpType::LoadFunction
-/// [Module]: hugr_core::ops::OpType::Module
 pub fn remove_dead_funcs(
     h: &mut impl HugrMut<Node = Node>,
-    entry_points: impl IntoIterator<Item = Node>,
 ) -> Result<(), ValidatePassError<Node, RemoveDeadFuncsError>> {
-    validate_if_test(
-        RemoveDeadFuncsPass::default().with_module_entry_points(entry_points),
-        h,
-    )
+    validate_if_test(RemoveDeadFuncsPass::default(), h)
 }
 
 #[cfg(test)]
@@ -146,29 +157,34 @@ mod test {
     };
     use hugr_core::{extension::prelude::usize_t, types::Signature, HugrView};
 
-    use super::remove_dead_funcs;
+    use super::RemoveDeadFuncsPass;
+    use crate::ComposablePass;
 
     #[rstest]
-    #[case([], vec![])] // No entry_points removes everything!
-    #[case(["main"], vec!["from_main", "main"])]
-    #[case(["from_main"], vec!["from_main"])]
-    #[case(["other1"], vec!["other1", "other2"])]
-    #[case(["other2"], vec!["other2"])]
-    #[case(["other1", "other2"], vec!["other1", "other2"])]
+    #[case(false, [], vec![])] // No entry_points removes everything!
+    #[case(false, ["main"], vec!["from_main", "main"])]
+    #[case(false, ["from_main"], vec!["from_main"])]
+    #[case(false, ["other1"], vec!["other1", "other2"])]
+    #[case(false, ["other2"], vec!["other2"])]
+    #[case(false, ["other1", "other2"], vec!["other1", "other2"])]
+    #[case(true, [], vec!["from_main", "main", "other2"])]
+    #[case(true, ["other1"], vec!["from_main", "main", "other1", "other2"])]
     fn remove_dead_funcs_entry_points(
+        #[case] include_exports: bool,
         #[case] entry_points: impl IntoIterator<Item = &'static str>,
         #[case] retained_funcs: Vec<&'static str>,
     ) -> Result<(), Box<dyn std::error::Error>> {
         let mut hb = ModuleBuilder::new();
         let o2 = hb.define_function("other2", Signature::new_endo(usize_t()))?;
         let o2inp = o2.input_wires();
         let o2 = o2.finish_with_outputs(o2inp)?;
-        let mut o1 = hb.define_function("other1", Signature::new_endo(usize_t()))?;
+        let mut o1 =
+            hb.define_function_link_name("other1", Signature::new_endo(usize_t()), None)?;
 
         let o1c = o1.call(o2.handle(), &[], o1.input_wires())?;
         o1.finish_with_outputs(o1c.outputs())?;
 
-        let fm = hb.define_function("from_main", Signature::new_endo(usize_t()))?;
+        let fm = hb.define_function_link_name("from_main", Signature::new_endo(usize_t()), None)?;
         let f_inp = fm.input_wires();
         let fm = fm.finish_with_outputs(f_inp)?;
         let mut m = hb.define_function("main", Signature::new_endo(usize_t()))?;
@@ -186,14 +202,16 @@ mod test {
             })
             .collect::<HashMap<_, _>>();
 
-        remove_dead_funcs(
-            &mut hugr,
-            entry_points
-                .into_iter()
-                .map(|name| *avail_funcs.get(name).unwrap())
-                .collect::<Vec<_>>(),
-        )
-        .unwrap();
+        RemoveDeadFuncsPass::default()
+            .include_module_exports(include_exports)
+            .with_module_entry_points(
+                entry_points
+                    .into_iter()
+                    .map(|name| *avail_funcs.get(name).unwrap())
+                    .collect::<Vec<_>>(),
+            )
+            .run(&mut hugr)
+            .unwrap();
 
         let remaining_funcs = hugr
             .nodes()
diff --git a/hugr-passes/src/monomorphize.rs b/hugr-passes/src/monomorphize.rs
@@ -307,7 +307,7 @@ mod test {
     use hugr_core::{Hugr, HugrView, Node};
     use rstest::rstest;
 
-    use crate::{monomorphize, remove_dead_funcs};
+    use crate::{monomorphize, remove_dead_funcs, ComposablePass, RemoveDeadFuncsPass};
 
     use super::{is_polymorphic, mangle_inner_func, mangle_name};
 
@@ -414,7 +414,10 @@ mod test {
         assert_eq!(mono2, mono); // Idempotent
 
         let mut nopoly = mono;
-        remove_dead_funcs(&mut nopoly, [mn.node()])?;
+        RemoveDeadFuncsPass::default()
+            .include_module_exports(false)
+            .with_module_entry_points([mn.node()])
+            .run(&mut nopoly)?;
         let mut funcs = list_funcs_link_name(&nopoly);
 
         assert!(funcs.values().all(|(_, fd)| !is_polymorphic(fd)));
@@ -617,12 +620,13 @@ mod test {
             let mut module_builder = ModuleBuilder::new();
             let foo = {
                 let builder = module_builder
-                    .define_function(
+                    .define_function_link_name(
                         "foo",
                         PolyFuncType::new(
                             [TypeBound::Any.into()],
                             Signature::new_endo(Type::new_var_use(0, TypeBound::Any)),
                         ),
+                        None,
                     )
                     .unwrap();
                 let inputs = builder.input_wires();
@@ -653,7 +657,7 @@ mod test {
         };
 
         monomorphize(&mut hugr).unwrap();
-        remove_dead_funcs(&mut hugr, []).unwrap();
+        remove_dead_funcs(&mut hugr).unwrap();
 
         let funcs = list_funcs(&hugr);
         assert!(funcs.values().all(|(_, fd)| !is_polymorphic(fd)));
