feat: update passes to use PassScope where non-breaking

hugr-passes/src/composable.rs

@@ -307,6 +307,7 @@ impl<
     }
 }
 
+// Note remove when deprecated constant_fold_pass / remove_dead_funcs are removed
 pub(crate) fn validate_if_test<P: ComposablePass<H>, H: HugrMut>(
     pass: P,
     hugr: &mut H,
@@ -319,7 +320,7 @@ pub(crate) fn validate_if_test<P: ComposablePass<H>, H: HugrMut>(
 }
 
 #[cfg(test)]
-mod test {
+pub(crate) mod test {
     use hugr_core::ops::Value;
     use itertools::{Either, Itertools};
 
@@ -335,11 +336,18 @@ mod test {
 
     use crate::const_fold::{ConstFoldError, ConstantFoldPass};
     use crate::dead_code::DeadCodeElimError;
-    use crate::untuple::{UntupleRecursive, UntupleResult};
-    use crate::{DeadCodeElimPass, ReplaceTypes, UntuplePass};
+    use crate::untuple::UntupleResult;
+    use crate::{DeadCodeElimPass, PassScope, ReplaceTypes, UntuplePass};
 
     use super::{ComposablePass, IfThen, ValidatePassError, ValidatingPass, validate_if_test};
 
+    pub(crate) fn run_validating<P: ComposablePass<H>, H: HugrMut>(
+        pass: P,
+        hugr: &mut H,
+    ) -> Result<P::Result, ValidatePassError<H::Node, P::Error>> {
+        ValidatingPass::new(pass).run(hugr)
+    }
+
     #[test]
     fn test_then() {
         let mut mb = ModuleBuilder::new();
@@ -444,7 +452,7 @@ mod test {
             fb.finish_hugr_with_outputs(untup.outputs()).unwrap()
         };
 
-        let untup = UntuplePass::new(UntupleRecursive::Recursive);
+        let untup = UntuplePass::new_scoped(PassScope::EntrypointRecursive);
         {
             // Change usize_t to INT_TYPES[6], and if that did anything (it will!), then Untuple
             let mut repl = ReplaceTypes::default();
@@ -453,7 +461,7 @@ mod test {
             let ifthen = IfThen::<Either<_, _>, _, _, _>::new(repl, untup.clone());
 
             let mut h = h.clone();
-            let r = validate_if_test(ifthen, &mut h).unwrap();
+            let r = run_validating(ifthen, &mut h).unwrap();
             assert_eq!(
                 r,
                 Some(UntupleResult {
@@ -470,7 +478,7 @@ mod test {
         repl.set_replace_type(i32_custom_t, INT_TYPES[6].clone());
         let ifthen = IfThen::<Either<_, _>, _, _, _>::new(repl, untup);
         let mut h = h;
-        let r = validate_if_test(ifthen, &mut h).unwrap();
+        let r = run_validating(ifthen, &mut h).unwrap();
         assert_eq!(r, None);
         assert_eq!(h.children(h.entrypoint()).count(), 4);
         let mktup = h

hugr-passes/src/const_fold.rs

@@ -20,15 +20,16 @@ use crate::dataflow::{
     partial_from_const,
 };
 use crate::dead_code::{DeadCodeElimError, DeadCodeElimPass, PreserveNode};
-use crate::{ComposablePass, composable::validate_if_test};
+use crate::{ComposablePass, PassScope, composable::validate_if_test};
 
 #[derive(Debug, Clone, Default)]
 /// A configuration for the Constant Folding pass.
 pub struct ConstantFoldPass {
     allow_increase_termination: bool,
+    scope: Option<PassScope>,
     /// 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>>,
 }
 
@@ -67,8 +68,8 @@ impl ConstantFoldPass {
     }
 
     /// Specifies a number of external inputs to an entry point of the Hugr.
-    /// In normal use, for Module-rooted Hugrs, `node` is a `FuncDefn` child of the root;
-    /// or for non-Module-rooted Hugrs, `node` is the root of the Hugr. (This is not
+    /// In normal use, for Module-entrypoint Hugrs, `node` is a `FuncDefn` child of the module;
+    /// or for non-Module-entrypoint Hugrs, `node` is the entrypoint of the Hugr. (This is not
     /// enforced, but it must be a container and not a module itself.)
     ///
     /// Multiple calls for the same entry-point combine their values, with later
@@ -100,6 +101,13 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
     /// [`ConstFoldError::InvalidEntryPoint`] if an entry-point added by [`Self::with_inputs`]
     /// was of an invalid [`OpType`]
     fn run(&self, hugr: &mut H) -> Result<(), ConstFoldError> {
+        let Some(root) = self
+            .scope
+            .as_ref()
+            .map_or(Some(hugr.entrypoint()), |sc| sc.root(hugr))
+        else {
+            return Ok(()); // Scope says do nothing
+        };
         let fresh_node = Node::from(portgraph::NodeIndex::new(
             hugr.nodes().max().map_or(0, |n| n.index() + 1),
         ));
@@ -122,15 +130,24 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
             .map_err(|op| ConstFoldError::InvalidEntryPoint { node: n, op })?;
         }
 
+        for node in self.scope.iter().flat_map(|sc| sc.preserve_interface(hugr)) {
+            if node == hugr.module_root() || self.inputs.contains_key(&node) {
+                // Cannot prepopulate inputs for module-root; do not `join` with inputs explicitly specified.
+                continue;
+            }
+            const NO_INPUTS: [(IncomingPort, PartialValue<ValueHandle>); 0] = [];
+            m.prepopulate_inputs(node, NO_INPUTS)
+                .map_err(|op| ConstFoldError::InvalidEntryPoint { node, op })?;
+        }
+
         let results = m.run(ConstFoldContext, []);
         let mb_root_inp = hugr.get_io(hugr.entrypoint()).map(|[i, _]| i);
 
         let wires_to_break = hugr
-            .entry_descendants()
+            .descendants(root)
             .flat_map(|n| hugr.node_inputs(n).map(move |ip| (n, ip)))
             .filter(|(n, ip)| {
-                *n != hugr.entrypoint()
-                    && matches!(hugr.get_optype(*n).port_kind(*ip), Some(EdgeKind::Value(_)))
+                *n != root && matches!(hugr.get_optype(*n).port_kind(*ip), Some(EdgeKind::Value(_)))
             })
             .filter_map(|(n, ip)| {
                 let (src, outp) = hugr.single_linked_output(n, ip).unwrap();
@@ -165,8 +182,13 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
             hugr.connect(lcst, OutgoingPort::from(0), n, inport);
         }
         // Eliminate dead code not required for the same entry points.
-        DeadCodeElimPass::<H>::default()
-            .with_entry_points(self.inputs.keys().copied())
+        let dce = self
+            .scope
+            .as_ref()
+            .map_or(DeadCodeElimPass::<H>::default(), |scope| {
+                DeadCodeElimPass::<H>::default().with_scope_internal(scope.clone())
+            });
+        dce.with_entry_points(self.inputs.keys().copied())
             .set_preserve_callback(if self.allow_increase_termination {
                 Arc::new(|_, _| PreserveNode::CanRemoveIgnoringChildren)
             } else {
@@ -186,13 +208,19 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
             })?;
         Ok(())
     }
+
+    fn with_scope_internal(mut self, scope: impl Into<PassScope>) -> Self {
+        self.scope = Some(scope.into());
+        self
+    }
 }
 
 /// 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 ConstantFoldPass with a PassScope", since = "0.25.7")]
 pub fn constant_fold_pass<H: HugrMut<Node = Node> + 'static>(mut h: impl AsMut<H>) {
     let h = h.as_mut();
     let c = ConstantFoldPass::default();

hugr-passes/src/const_fold/test.rs

@@ -3,8 +3,8 @@ use std::{
     sync::LazyLock,
 };
 
-use hugr_core::ops::Const;
 use hugr_core::ops::handle::NodeHandle;
+use hugr_core::{Visibility, ops::Const};
 use itertools::Itertools;
 use rstest::rstest;
 
@@ -30,10 +30,20 @@ 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 _, composable::ValidatingPass};
+use crate::{
+    PassScope,
+    composable::WithScope,
+    dataflow::{DFContext, PartialValue, partial_from_const},
+};
+
+use super::{ConstFoldContext, ConstantFoldPass, ValueHandle};
 
-use super::{ConstFoldContext, ConstantFoldPass, ValueHandle, constant_fold_pass};
+fn constant_fold_pass(h: &mut (impl HugrMut<Node = Node> + 'static)) {
+    // the default ConstantFoldPass has no scope, i.e. preserving legacy behavior
+    let c = ConstantFoldPass::default().with_scope(PassScope::default());
+    ValidatingPass::new(c).run(h).unwrap();
+}
 
 #[rstest]
 #[case(ConstInt::new_u(4, 2).unwrap(), true)]
@@ -1592,9 +1602,10 @@ fn test_module() -> Result<(), Box<dyn std::error::Error>> {
     let c17 = mb.add_constant(Value::from(ConstInt::new_u(5, 17)?));
     let ad1 = mb.add_alias_declare("unused", TypeBound::Linear)?;
     let ad2 = mb.add_alias_def("unused2", INT_TYPES[3].clone())?;
-    let mut main = mb.define_function(
+    let mut main = mb.define_function_vis(
         "main",
         Signature::new(type_row![], vec![INT_TYPES[5].clone(); 2]),
+        Visibility::Public,
     )?;
     let lc7 = main.load_const(&c7);
     let lc17 = main.load_const(&c17);

hugr-passes/src/dead_code.rs

@@ -6,14 +6,16 @@ use std::collections::{HashMap, HashSet, VecDeque};
 use std::fmt::{Debug, Display, Formatter};
 use std::sync::Arc;
 
-use crate::ComposablePass;
+use crate::{ComposablePass, PassScope};
 
 /// Configuration for Dead Code Elimination pass
 #[derive(Clone)]
 pub struct DeadCodeElimPass<H: HugrView> {
     /// Nodes that are definitely needed - e.g. `FuncDefns`, but could be anything.
     /// Hugr Root is assumed to be an entry point even if not mentioned here.
     entry_points: Vec<H::Node>,
+    /// If None, use entrypoint-subtree (even if module root)
+    scope: Option<PassScope>,
     /// Callback identifying nodes that must be preserved even if their
     /// results are not used. Defaults to [`PreserveNode::default_for`].
     preserve_callback: Arc<PreserveCallback<H>>,
@@ -23,6 +25,8 @@ impl<H: HugrView + 'static> Default for DeadCodeElimPass<H> {
     fn default() -> Self {
         Self {
             entry_points: Default::default(),
+            // Preserve pre-PassScope behaviour of affecting entrypoint subtree only:
+            scope: None,
             preserve_callback: Arc::new(PreserveNode::default_for),
         }
     }
@@ -36,11 +40,13 @@ impl<H: HugrView> Debug for DeadCodeElimPass<H> {
         #[derive(Debug)]
         struct DCEDebug<'a, N> {
             entry_points: &'a Vec<N>,
+            scope: &'a Option<PassScope>,
         }
 
         Debug::fmt(
             &DCEDebug {
                 entry_points: &self.entry_points,
+                scope: &self.scope,
             },
             f,
         )
@@ -97,11 +103,11 @@ impl<H: HugrView> DeadCodeElimPass<H> {
         self
     }
 
-    /// 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.
+    /// Mark some nodes as starting points for analysis, i.e. so we cannot eliminate any code
+    /// used to evaluate these nodes. (E.g. nodes at which we may start executing the Hugr.)
+    ///
+    /// Other starting points are added according to the [PassScope].
+    // TODO should we deprecate this? i.e. require use of PreserveCallback / Hugr edges?
     pub fn with_entry_points(mut self, entry_points: impl IntoIterator<Item = H::Node>) -> Self {
         self.entry_points.extend(entry_points);
         self
@@ -111,14 +117,21 @@ impl<H: HugrView> DeadCodeElimPass<H> {
         let mut must_preserve = HashMap::new();
         let mut needed = HashSet::new();
         let mut q = VecDeque::from_iter(self.entry_points.iter().copied());
-        q.push_front(h.entrypoint());
+
+        match &self.scope {
+            None => q.push_back(h.entrypoint()),
+            Some(scope) => q.extend(scope.preserve_interface(h)),
+        };
         while let Some(n) = q.pop_front() {
             if !h.contains_node(n) {
                 return Err(DeadCodeElimError::NodeNotFound(n));
             }
             if !needed.insert(n) {
                 continue;
             }
+            // Ensure no orphans. We could remove more from parent, but would require transforming
+            // (e.g. removing individual Output ports) not just deleting, so don't.
+            q.extend(h.get_parent(n));
             for (i, ch) in h.children(n).enumerate() {
                 if self.must_preserve(h, &mut must_preserve, ch)
                     || match h.get_optype(ch) {
@@ -181,16 +194,28 @@ impl<H: HugrMut> ComposablePass<H> for DeadCodeElimPass<H> {
     type Result = ();
 
     fn run(&self, hugr: &mut H) -> Result<(), Self::Error> {
+        let root = match &self.scope {
+            None => hugr.entrypoint(),
+            Some(scope) => match scope.root(hugr) {
+                Some(root) => root,
+                None => return Ok(()),
+            },
+        };
         let needed = self.find_needed_nodes(&*hugr)?;
         let remove = hugr
-            .entry_descendants()
+            .descendants(root)
             .filter(|n| !needed.contains(n))
             .collect::<Vec<_>>();
         for n in remove {
             hugr.remove_node(n);
         }
         Ok(())
     }
+
+    fn with_scope_internal(mut self, scope: impl Into<PassScope>) -> Self {
+        self.scope = Some(scope.into());
+        self
+    }
 }
 #[cfg(test)]
 mod test {