Revert "Revert hugr-passes changes"

This reverts commit 3a460a4.

Author: acl-cqc

hugr-passes/src/call_graph.rs

@@ -18,7 +18,7 @@ pub enum CallGraphNode<N = Node> {
     FuncDecl(N),
     /// petgraph-node corresponds to a [`FuncDefn`](OpType::FuncDefn) node (specified) in the Hugr
     FuncDefn(N),
-    /// petgraph-node corresponds to the root node of the hugr, that is not
+    /// petgraph-node corresponds to the entrypoint node of the hugr, that is not
     /// a [`FuncDefn`](OpType::FuncDefn). Note that it will not be a [Module](OpType::Module)
     /// either, as such a node could not have outgoing edges, so is not represented in the petgraph.
     NonFuncRoot,

hugr-passes/src/const_fold.rs

@@ -15,38 +15,44 @@ 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.
+///
+/// Note that by default we assume that only the entrypoint is reachable and
+/// only if it is not the module root; see [Self::with_inputs]. Mutation
+/// occurs anywhere beneath the entrypoint.
 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
+    ///   - a `FuncDefn` child of the module-root
+    ///   - the entrypoint
     inputs: HashMap<Node, HashMap<IncomingPort, Value>>,
 }
 
 #[derive(Clone, Debug, Error, PartialEq)]
 #[non_exhaustive]
 /// Errors produced by [`ConstantFoldPass`].
 pub enum ConstFoldError {
-    /// Error raised when a Node is specified as an entry-point but
-    /// is neither a dataflow parent, nor a [CFG](OpType::CFG), nor
-    /// a [Conditional](OpType::Conditional).
+    /// Error raised when inputs are provided for a Node that is neither a dataflow
+    /// parent, nor a [CFG](OpType::CFG), nor a [Conditional](OpType::Conditional).
     #[error("{node} has OpType {op} which cannot be an entry-point")]
     InvalidEntryPoint {
         /// The node which was specified as an entry-point
         node: Node,
         /// The `OpType` of the node
         op: OpType,
     },
-    /// The chosen entrypoint is not in the hugr.
+    /// Inputs were provided for a node that is not in the hugr.
     #[error("Entry-point {node} is not part of the Hugr")]
     MissingEntryPoint {
         /// The missing node
@@ -67,15 +73,25 @@ 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-rooted Hugrs, `node` is a `FuncDefn` (child of the root);
+    /// for non-Module-rooted Hugrs, `node` is the [HugrView::entrypoint]. (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
     /// values on the same in-port replacing earlier ones.
     ///
-    /// Note that if `inputs` is empty, this still marks the node as an entry-point, i.e.
-    /// we must preserve nodes required to compute its result.
+    /// Note that providing empty `inputs` indicates that we must preserve the ability
+    /// to compute the result of `node` for all possible inputs.
+    /// * If the entrypoint is the module-root, this method should be called for every
+    /// [FuncDefn] that is externally callable
+    /// * Otherwise, i.e. if the entrypoint is not the module-root,
+    ///    * The default is to assume the entrypoint is callable with any inputs;
+    ///    * If `node` is the entrypoint, this method allows to restrict the possible inputs
+    ///    * If `node` is beneath the entrypoint, this merely degrades the analysis. (We
+    ///      will mutate only beneath the entrypoint, but using results of analysing the
+    ///      whole Hugr wrt. the specified/any inputs too).
+    /// 
+    /// [FuncDefn]: hugr_core::ops::FuncDefn
     pub fn with_inputs(
         mut self,
         node: Node,
@@ -97,8 +113,7 @@ impl<H: HugrMut<Node = Node> + 'static> ComposablePass<H> for ConstantFoldPass {
     ///
     /// # Errors
     ///
-    /// [`ConstFoldError::InvalidEntryPoint`] if an entry-point added by [`Self::with_inputs`]
-    /// was of an invalid [`OpType`]
+    /// [ConstFoldError] if inputs were provided via [`Self::with_inputs`] for an invalid node.
     fn run(&self, hugr: &mut H) -> Result<(), ConstFoldError> {
         let fresh_node = Node::from(portgraph::NodeIndex::new(
             hugr.nodes().max().map_or(0, |n| n.index() + 1),
@@ -184,25 +199,51 @@ 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.
+/// Otherwise, assume that the [HugrView::entrypoint] is itself reachable.
 ///
 /// [`FuncDefn`]: hugr_core::ops::OpType::FuncDefn
 /// [`Module`]: hugr_core::ops::OpType::Module
+#[deprecated(note = "Use fold_constants, 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 fold_constants(h: &mut (impl HugrMut<Node = Node> + 'static), policy: IncludeExports) {
+    let mut funcs = Vec::new();
+    if !h.entrypoint_optype().is_module() {
+        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, fold_constants};
+
+fn constant_fold_pass(h: &mut (impl HugrMut<Node = Node> + 'static)) {
+    fold_constants(h, IncludeExports::Always);
+}
 
 #[rstest]
 #[case(ConstInt::new_u(4, 2).unwrap(), true)]

hugr-passes/src/dataflow/datalog.rs

@@ -116,7 +116,7 @@ impl<H: HugrView, V: AbstractValue> Machine<H, V> {
         } else {
             let ep = self.0.entrypoint();
             let mut p = in_values.into_iter().peekable();
-            // We must provide some inputs to the root so that they are Top rather than Bottom.
+            // We must provide some inputs to the entrypoint so that they are Top rather than Bottom.
             // (However, this test will fail for DataflowBlock or Case roots, i.e. if no
             // inputs have been provided they will still see Bottom. We could store the "input"
             // values for even these nodes in self.1 and then convert to actual Wire values

hugr-passes/src/dead_code.rs

@@ -1,32 +1,35 @@
 //! Pass for removing dead code, i.e. that computes values that are then discarded
 
 use hugr_core::hugr::internal::HugrInternals;
-use hugr_core::{HugrView, hugr::hugrmut::HugrMut, ops::OpType};
+use hugr_core::{HugrView, Visibility, hugr::hugrmut::HugrMut, ops::OpType};
 use std::convert::Infallible;
 use std::fmt::{Debug, Formatter};
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     sync::Arc,
 };
 
-use crate::ComposablePass;
+use crate::{ComposablePass, IncludeExports};
 
-/// Configuration for Dead Code Elimination pass
+/// Configuration for Dead Code Elimination pass, i.e. which removes nodes
+/// beneath the [HugrView::entrypoint] that compute only unneeded values.
 #[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.
+    /// [HugrView::entrypoint] is assumed to be needed even if not mentioned here.
     entry_points: Vec<H::Node>,
     /// 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: IncludeExports,
 }
 
 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: IncludeExports::default(),
         }
     }
 }
@@ -39,11 +42,13 @@ impl<H: HugrView> Debug for DeadCodeElimPass<H> {
         #[derive(Debug)]
         struct DCEDebug<'a, N> {
             entry_points: &'a Vec<N>,
+            include_exports: IncludeExports,
         }
 
         Debug::fmt(
             &DCEDebug {
                 entry_points: &self.entry_points,
+                include_exports: self.include_exports,
             },
             f,
         )
@@ -69,12 +74,12 @@ pub enum PreserveNode {
 
 impl PreserveNode {
     /// A conservative default for a given node. Just examines the node's [`OpType`]:
-    /// * Assumes all Calls must be preserved. (One could scan the called `FuncDefn`, but would
-    ///   also need to check for cycles in the [`CallGraph`](super::call_graph::CallGraph).)
+    /// * Assumes all Calls must be preserved. (One could scan the called `FuncDefn` for
+    ///   termination, but would also need to check for cycles in the `CallGraph`.)
     /// * Assumes all CFGs must be preserved. (One could, for example, allow acyclic
     ///   CFGs to be removed.)
-    /// * Assumes all `TailLoops` must be preserved. (One could, for example, use dataflow
-    ///   analysis to allow removal of `TailLoops` that never [Continue](hugr_core::ops::TailLoop::CONTINUE_TAG).)
+    /// * Assumes all `TailLoops` must be preserved. (One could use some analysis, e.g.
+    ///   dataflow, to allow removal of `TailLoops` with a bounded number of iterations.)
     pub fn default_for<H: HugrView>(h: &H, n: H::Node) -> PreserveNode {
         match h.get_optype(n) {
             OpType::CFG(_) | OpType::TailLoop(_) | OpType::Call(_) => PreserveNode::MustKeep,
@@ -91,16 +96,33 @@ 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 reachable, i.e. so we cannot eliminate any code used to
+    /// evaluate their results. The [`HugrView::entrypoint`] is assumed to be reachable;
+    /// if that is the [`HugrView::module_root`], then any public [FuncDefn] and
+    /// [FuncDecl]s are also considered reachable by default,
+    /// but this can be change by [`Self::include_module_exports`].
+    ///
+    /// [FuncDecl]: OpType::FuncDecl
+    /// [FuncDefn]: OpType::FuncDefn
     pub fn with_entry_points(mut self, entry_points: impl IntoIterator<Item = H::Node>) -> Self {
         self.entry_points.extend(entry_points);
         self
     }
 
+    /// Sets whether the exported [FuncDefn](OpType::FuncDefn)s and
+    /// [FuncDecl](OpType::FuncDecl)s are considered reachable.
+    ///
+    /// Note that for non-module-entry Hugrs this has no effect, since we only remove
+    /// code beneath the entrypoint: this cannot be affected by other module children.
+    ///
+    /// So, for module-rooted-Hugrs: [IncludeExports::OnlyIfEntrypointIsModuleRoot] is
+    /// equivalent to [IncludeExports::Always]; and [IncludeExports::Never] will remove
+    /// all children, unless some are explicity added by [Self::with_entry_points].
+    pub fn include_module_exports(mut self, include: IncludeExports) -> 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();
@@ -111,19 +133,23 @@ impl<H: HugrView> DeadCodeElimPass<H> {
                 continue;
             }
             for ch in h.children(n) {
-                if self.must_preserve(h, &mut must_preserve, ch)
-                    || matches!(
-                        h.get_optype(ch),
+                let must_keep = match h.get_optype(ch) {
                         OpType::Case(_) // Include all Cases in Conditionals
                     | OpType::DataflowBlock(_) // and all Basic Blocks in CFGs
                     | OpType::ExitBlock(_)
                     | 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(_) => true,
+                    // FuncDefns (as children of Module) only if public and including exports
+                    // (will be included if static predecessors of Call/LoadFunction below,
+                    // regardless of Visibility or self.include_exports)
+                    OpType::FuncDefn(fd) => fd.visibility() == &Visibility::Public && self.include_exports.for_hugr(h),
+                    OpType::FuncDecl(fd) => fd.visibility() == &Visibility::Public && self.include_exports.for_hugr(h),
+                    // No Const, unless reached along static edges
+                    _ => false
+                };
+                if must_keep || self.must_preserve(h, &mut must_preserve, ch) {
                     q.push_back(ch);
                 }
             }
@@ -141,7 +167,6 @@ impl<H: HugrView> DeadCodeElimPass<H> {
         if let Some(res) = cache.get(&n) {
             return *res;
         }
-        #[allow(deprecated)]
         let res = match self.preserve_callback.as_ref()(h, n) {
             PreserveNode::MustKeep => true,
             PreserveNode::CanRemoveIgnoringChildren => false,
@@ -174,18 +199,57 @@ impl<H: HugrMut> ComposablePass<H> for DeadCodeElimPass<H> {
 mod test {
     use std::sync::Arc;
 
-    use hugr_core::Hugr;
-    use hugr_core::builder::{CFGBuilder, Container, Dataflow, DataflowSubContainer, HugrBuilder};
+    use hugr_core::builder::{
+        CFGBuilder, Container, Dataflow, DataflowSubContainer, HugrBuilder, ModuleBuilder,
+    };
     use hugr_core::extension::prelude::{ConstUsize, usize_t};
-    use hugr_core::ops::{OpTag, OpTrait, handle::NodeHandle};
-    use hugr_core::types::Signature;
-    use hugr_core::{HugrView, ops::Value, type_row};
+    use hugr_core::ops::{OpTag, OpTrait, Value, handle::NodeHandle};
+    use hugr_core::{Hugr, HugrView, type_row, types::Signature};
     use itertools::Itertools;
+    use rstest::rstest;
 
-    use crate::ComposablePass;
+    use crate::{ComposablePass, IncludeExports};
 
     use super::{DeadCodeElimPass, PreserveNode};
 
+    #[rstest]
+    #[case(false, IncludeExports::Never, true)]
+    #[case(false, IncludeExports::OnlyIfEntrypointIsModuleRoot, false)]
+    #[case(false, IncludeExports::Always, false)]
+    #[case(true, IncludeExports::Never, true)]
+    #[case(true, IncludeExports::OnlyIfEntrypointIsModuleRoot, false)]
+    #[case(true, IncludeExports::Always, false)]
+    fn test_module_exports(
+        #[case] include_dfn: bool,
+        #[case] module_exports: IncludeExports,
+        #[case] decl_removed: bool,
+    ) {
+        let mut mb = ModuleBuilder::new();
+        let dfn = mb
+            .define_function("foo", Signature::new_endo(usize_t()))
+            .unwrap();
+        let ins = dfn.input_wires();
+        let dfn = dfn.finish_with_outputs(ins).unwrap();
+        let dcl = mb
+            .declare("bar", Signature::new_endo(usize_t()).into())
+            .unwrap();
+        let mut h = mb.finish_hugr().unwrap();
+        let mut dce = DeadCodeElimPass::<Hugr>::default().include_module_exports(module_exports);
+        if include_dfn {
+            dce = dce.with_entry_points([dfn.node()]);
+        }
+        dce.run(&mut h).unwrap();
+        let defn_retained = include_dfn;
+        let decl_retained = !decl_removed;
+        let children = h.children(h.module_root()).collect_vec();
+        assert_eq!(defn_retained, children.iter().contains(&dfn.node()));
+        assert_eq!(decl_retained, children.iter().contains(&dcl.node()));
+        assert_eq!(
+            children.len(),
+            (defn_retained as usize) + (decl_retained as usize)
+        );
+    }
+
     #[test]
     fn test_cfg_callback() {
         let mut cb = CFGBuilder::new(Signature::new_endo(type_row![])).unwrap();