Move IncludeExports up to root of hugr-passes, use for const_fold, +deprecate

Author: acl-cqc

hugr-passes/src/const_fold.rs

@@ -16,20 +16,23 @@ use hugr_core::{
 };
 use value_handle::ValueHandle;
 
-use crate::dataflow::{
-    ConstLoader, ConstLocation, DFContext, Machine, PartialValue, TailLoopTermination,
-    partial_from_const,
-};
 use crate::dead_code::{DeadCodeElimPass, PreserveNode};
 use crate::{ComposablePass, composable::validate_if_test};
+use crate::{
+    IncludeExports,
+    dataflow::{
+        ConstLoader, ConstLocation, DFContext, Machine, PartialValue, TailLoopTermination,
+        partial_from_const,
+    },
+};
 
 #[derive(Debug, Clone, Default)]
 /// A configuration for the Constant Folding pass.
 pub struct ConstantFoldPass {
     allow_increase_termination: bool,
     /// Each outer key Node must be either:
     ///   - a `FuncDefn` child of the root, if the root is a module; or
-    ///   - the root, if the root is not a Module
+    ///   - the entrypoint, if the entrypoint is not a Module
     inputs: HashMap<Node, HashMap<IncomingPort, Value>>,
 }
 
@@ -185,25 +188,53 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
     }
 }
 
+const NO_INPUTS: [(IncomingPort, Value); 0] = [];
+
 /// Exhaustively apply constant folding to a HUGR.
 /// If the Hugr's entrypoint is its [`Module`], assumes all [`FuncDefn`] children are reachable.
 ///
 /// [`FuncDefn`]: hugr_core::ops::OpType::FuncDefn
 /// [`Module`]: hugr_core::ops::OpType::Module
+#[deprecated(note = "Use constant_fold_pass_pub, or manually configure ConstantFoldPass")]
 pub fn constant_fold_pass<H: HugrMut<Node = Node> + 'static>(mut h: impl AsMut<H>) {
     let h = h.as_mut();
     let c = ConstantFoldPass::default();
     let c = if h.get_optype(h.entrypoint()).is_module() {
-        let no_inputs: [(IncomingPort, _); 0] = [];
         h.children(h.entrypoint())
             .filter(|n| h.get_optype(*n).is_func_defn())
-            .fold(c, |c, n| c.with_inputs(n, no_inputs.iter().cloned()))
+            .fold(c, |c, n| c.with_inputs(n, NO_INPUTS.clone()))
     } else {
         c
     };
     validate_if_test(c, h).unwrap();
 }
 
+/// Exhaustively apply constant folding to a HUGR.
+/// Assumes that the Hugr's entrypoint is reachable (if it is not a [`Module`]).
+/// Also uses `policy` to determine which public [`FuncDefn`] children of the [`HugrView::module_root`] are reachable.
+///
+/// [`Module`]: hugr_core::ops::OpType::Module
+/// [`FuncDefn`]: hugr_core::ops::OpType::FuncDefn
+pub fn constant_fold_pass_pub(
+    h: &mut (impl HugrMut<Node = Node> + 'static),
+    policy: IncludeExports,
+) {
+    let mut funcs = Vec::new();
+    if h.get_optype(h.entrypoint()).is_func_defn() {
+        funcs.push(h.entrypoint());
+    }
+    if policy.for_hugr(&h) {
+        funcs.extend(
+            h.children(h.module_root())
+                .filter(|n| h.get_optype(*n).is_func_defn()),
+        )
+    }
+    let c = funcs.into_iter().fold(ConstantFoldPass::default(), |c, n| {
+        c.with_inputs(n, NO_INPUTS.clone())
+    });
+    validate_if_test(c, h).unwrap();
+}
+
 struct ConstFoldContext;
 
 impl ConstLoader<ValueHandle<Node>> for ConstFoldContext {

hugr-passes/src/const_fold/test.rs

@@ -29,10 +29,14 @@ use hugr_core::std_extensions::logic::LogicOp;
 use hugr_core::types::{Signature, SumType, Type, TypeBound, TypeRow, TypeRowRV};
 use hugr_core::{Hugr, HugrView, IncomingPort, Node, type_row};
 
-use crate::ComposablePass as _;
 use crate::dataflow::{DFContext, PartialValue, partial_from_const};
+use crate::{ComposablePass as _, IncludeExports};
 
-use super::{ConstFoldContext, ConstantFoldPass, ValueHandle, constant_fold_pass};
+use super::{ConstFoldContext, ConstantFoldPass, ValueHandle, constant_fold_pass_pub};
+
+fn constant_fold_pass(h: &mut (impl HugrMut<Node = Node> + 'static)) {
+    constant_fold_pass_pub(h, IncludeExports::Always);
+}
 
 #[rstest]
 #[case(ConstInt::new_u(4, 2).unwrap(), true)]

hugr-passes/src/dead_funcs.rs

@@ -11,7 +11,7 @@ use hugr_core::{
 use petgraph::visit::{Dfs, Walker};
 
 use crate::{
-    ComposablePass,
+    ComposablePass, IncludeExports,
     composable::{ValidatePassError, validate_if_test},
 };
 
@@ -31,14 +31,6 @@ pub enum RemoveDeadFuncsError<N = Node> {
     },
 }
 
-#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
-pub enum IncludeExports {
-    Always,
-    Never,
-    #[default]
-    OnlyIfEntrypointIsModuleRoot,
-}
-
 fn reachable_funcs<'a, H: HugrView>(
     cg: &'a CallGraph<H::Node>,
     h: &'a H,
@@ -90,20 +82,7 @@ impl<H: HugrMut<Node = Node>> ComposablePass<H> for RemoveDeadFuncsPass {
     type Result = ();
     fn run(&self, hugr: &mut H) -> Result<(), RemoveDeadFuncsError> {
         let mut entry_points = Vec::new();
-        let include_exports = match self.include_exports {
-            IncludeExports::Always => true,
-            IncludeExports::Never => false,
-            IncludeExports::OnlyIfEntrypointIsModuleRoot => hugr.entrypoint() == hugr.module_root(),
-        };
-        let include_exports2 = matches!(
-            (
-                self.include_exports,
-                hugr.entrypoint() == hugr.module_root()
-            ),
-            (IncludeExports::Always, _) | (IncludeExports::OnlyIfEntrypointIsModuleRoot, true)
-        );
-        assert_eq!(include_exports, include_exports2);
-        if include_exports {
+        if self.include_exports.for_hugr(hugr) {
             entry_points.extend(hugr.children(hugr.module_root()).filter(|ch| {
                 hugr.get_optype(*ch)
                     .as_func_defn()

hugr-passes/src/lib.rs

@@ -13,6 +13,7 @@ pub use dead_funcs::{RemoveDeadFuncsError, RemoveDeadFuncsPass, remove_dead_func
 pub mod force_order;
 mod half_node;
 pub mod linearize_array;
+use hugr_core::HugrView;
 pub use linearize_array::LinearizeArrayPass;
 pub mod lower;
 pub mod merge_bbs;
@@ -28,3 +29,23 @@ pub use force_order::{force_order, force_order_by_key};
 pub use lower::{lower_ops, replace_many_ops};
 pub use non_local::{ensure_no_nonlocal_edges, nonlocal_edges};
 pub use untuple::UntuplePass;
+
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
+/// A policy for whether to include the public (exported) functions of a Hugr
+/// (typically, as starting points for analysis)
+pub enum IncludeExports {
+    Always,
+    Never,
+    #[default]
+    OnlyIfEntrypointIsModuleRoot,
+}
+
+impl IncludeExports {
+    /// Returns whether to include the public functions of a particular Hugr
+    fn for_hugr(&self, h: &impl HugrView) -> bool {
+        matches!(
+            (self, h.entrypoint() == h.module_root()),
+            (Self::Always, _) | (Self::OnlyIfEntrypointIsModuleRoot, true)
+        )
+    }
+}

hugr-passes/src/monomorphize.rs

@@ -290,7 +290,7 @@ mod test {
     use hugr_core::{Hugr, HugrView, Node};
     use rstest::rstest;
 
-    use crate::dead_funcs::IncludeExports;
+    use crate::IncludeExports;
     use crate::{ComposablePass, RemoveDeadFuncsPass, monomorphize, remove_dead_funcs};
 
     use super::{is_polymorphic, mangle_name};