Include in DeadCode/DeadFunc elimination

Author: acl-cqc

hugr-core/src/ops/constant.rs

@@ -56,6 +56,12 @@ impl Const {
         &self.value
     }
 
+    /// Whether the [`Const`] would be exported for linking
+    /// (only applies if the parent is a [Module](super::Module))
+    pub fn is_public(&self) -> bool {
+        self.name.is_some()
+    }
+
     delegate! {
         to self.value {
             /// Returns the type of this constant.

hugr-passes/src/dead_code.rs

@@ -19,6 +19,7 @@ pub struct DeadCodeElimPass {
     /// results are not used. Defaults to [PreserveNode::default_for].
     preserve_callback: Arc<PreserveCallback>,
     validation: ValidationLevel,
+    include_exports: bool,
 }
 
 impl Default for DeadCodeElimPass {
@@ -27,6 +28,7 @@ impl Default for DeadCodeElimPass {
             entry_points: Default::default(),
             preserve_callback: Arc::new(PreserveNode::default_for),
             validation: ValidationLevel::default(),
+            include_exports: true,
         }
     }
 }
@@ -102,18 +104,32 @@ impl DeadCodeElimPass {
 
     /// Mark some nodes as entry points to the Hugr, i.e. so we cannot eliminate any code
     /// used to evaluate these nodes.
-    /// The root node 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.
+    /// The root node is assumed to be an entry point; for Module roots, 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].
     pub fn with_entry_points(mut self, entry_points: impl IntoIterator<Item = 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: impl HugrView<Node = Node>) -> HashSet<Node> {
         let mut must_preserve = HashMap::new();
         let mut needed = HashSet::new();
         let mut q = VecDeque::from_iter(self.entry_points.iter().cloned());
+        if self.include_exports && h.root_type().is_module() {
+            q.extend(h.children(h.root()).filter(|ch| {
+                let op = h.get_optype(*ch);
+                op.as_func_defn().is_some_and(|fd| fd.public)
+                    || op.as_const().is_some_and(|c| c.is_public())
+            }))
+        }
         q.push_front(h.root());
         while let Some(n) = q.pop_front() {
             if !needed.insert(n) {

hugr-passes/src/dead_funcs.rs

@@ -61,11 +61,22 @@ fn reachable_funcs<'a, H: HugrView>(
     }))
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 /// A configuration for the Dead Function Removal pass.
 pub struct RemoveDeadFuncsPass {
     validation: ValidationLevel,
     entry_points: Vec<Node>,
+    include_exports: bool,
+}
+
+impl Default for RemoveDeadFuncsPass {
+    fn default() -> Self {
+        Self {
+            validation: Default::default(),
+            entry_points: Default::default(),
+            include_exports: true,
+        }
+    }
 }
 
 impl RemoveDeadFuncsPass {
@@ -88,10 +99,28 @@ impl RemoveDeadFuncsPass {
         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
+    }
+
     /// Runs the pass (see [remove_dead_funcs]) with this configuration
     pub fn run<H: HugrMut>(&self, hugr: &mut H) -> Result<(), RemoveDeadFuncsError> {
         self.validation.run_validated_pass(hugr, |hugr: &mut H, _| {
-            remove_dead_funcs(hugr, self.entry_points.iter().cloned())
+            let exports = if hugr.root_type().is_module() && self.include_exports {
+                hugr.children(hugr.root())
+                    .filter(|ch| {
+                        hugr.get_optype(*ch)
+                            .as_func_defn()
+                            .is_some_and(|fd| fd.public)
+                    })
+                    .collect()
+            } else {
+                vec![]
+            };
+            remove_dead_funcs(hugr, self.entry_points.iter().cloned().chain(exports))
         })
     }
 }
@@ -145,26 +174,29 @@ mod test {
     use super::RemoveDeadFuncsPass;
 
     #[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_vis("other1", Signature::new_endo(usize_t()), false)?;
 
         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_vis("from_main", Signature::new_endo(usize_t()), false)?;
         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()))?;
@@ -183,6 +215,7 @@ mod test {
             .collect::<HashMap<_, _>>();
 
         RemoveDeadFuncsPass::default()
+            .include_module_exports(include_exports)
             .with_module_entry_points(
                 entry_points
                     .into_iter()