feat: (WIP) Static evaluation via constant folding

hugr-core/src/hugr/rewrite/replace.rs

@@ -61,7 +61,7 @@ pub struct Replacement {
     /// and there would be no possible [Self::mu_inp], [Self::mu_out] or [Self::adoptions].
     pub removal: Vec<Node>,
     /// A hugr (not necessarily valid, as it may be missing edges and/or nodes), whose root
-    /// is the same type as the root of [Self::replacement].  "G" in the spec.
+    /// is the same type as the common parent of all the nodes in [Self::removal].  "G" in the spec.
     pub replacement: Hugr,
     /// Describes how parts of the Hugr that would otherwise be removed should instead be preserved but
     /// with new parents amongst the newly-inserted nodes.  This is a Map from container nodes in

hugr-core/src/types.rs

@@ -256,6 +256,15 @@ impl SumType {
             _ => None,
         }
     }
+
+    /// Convert the SumType into its constituent rows (one per variant)
+    pub fn into_rows(self) -> impl Iterator<Item = TypeRowRV> {
+        match self {
+            SumType::Unit { size } => vec![TypeRowRV::new(); size as usize],
+            SumType::General { rows } => rows,
+        }
+        .into_iter()
+    }
 }
 
 impl<RV: MaybeRV> From<SumType> for TypeBase<RV> {

hugr-passes/src/const_fold.rs

@@ -32,6 +32,7 @@ pub struct ConstantFoldPass {
     validation: ValidationLevel,
     allow_increase_termination: bool,
     inputs: HashMap<IncomingPort, Value>,
+    entry_point: Option<Node>,
 }
 
 #[derive(Debug, Error)]
@@ -71,6 +72,17 @@ impl ConstantFoldPass {
         self
     }
 
+    /// Sets the entry point for a [Module]-rooted Hugr, i.e. at which [FuncDefn]
+    /// child of the root the Hugr starts executing. If unspecified, for
+    /// [Module]-rooted Hugrs, the entire contents will be removed.
+    ///
+    /// [FuncDefn]: hugr_core::ops::OpType::FuncDefn
+    /// [Module]: hugr_core::ops::OpType::Module
+    pub fn with_module_entry_point(mut self, node: Node) -> Self {
+        self.entry_point = Some(node);
+        self
+    }
+
     /// Run the Constant Folding pass.
     fn run_no_validate(&self, hugr: &mut impl HugrMut) -> Result<(), ConstFoldError> {
         let fresh_node = Node::from(portgraph::NodeIndex::new(
@@ -88,8 +100,7 @@ impl ConstantFoldPass {
         });
 
         let results = Machine::new(&hugr).run(ConstFoldContext(hugr), inputs);
-        let mut keep_nodes = HashSet::new();
-        self.find_needed_nodes(&results, &mut keep_nodes);
+        let keep_nodes = self.find_needed_nodes(&results);
         let mb_root_inp = hugr.get_io(hugr.root()).map(|[i, _]| i);
 
         let remove_nodes = hugr
@@ -145,17 +156,23 @@ impl ConstantFoldPass {
     fn find_needed_nodes<H: HugrView>(
         &self,
         results: &AnalysisResults<ValueHandle, H>,
-        needed: &mut HashSet<Node>,
-    ) {
-        let mut q = VecDeque::new();
+    ) -> HashSet<Node> {
         let h = results.hugr();
-        q.push_back(h.root());
+        let mut needed = HashSet::new();
+        let mut q = VecDeque::from_iter([h.root()]);
+        if let Some(entry) = self.entry_point {
+            assert!(h.get_parent(entry) == Some(h.root()));
+            assert!(h.get_optype(entry).is_func_defn());
+            assert!(h.get_optype(h.root()).is_module());
+            q.push_back(entry);
+        }
+
         while let Some(n) = q.pop_front() {
             if !needed.insert(n) {
                 continue;
             };
 
-            if h.get_optype(n).is_cfg() {
+            if h.get_optype(n).is_cfg() || h.get_optype(n).is_conditional() {
                 for bb in h.children(n) {
                     //if results.bb_reachable(bb).unwrap() { // no, we'd need to patch up predicates
                     q.push_back(bb);
@@ -176,22 +193,19 @@ impl ConstantFoldPass {
             }
             // Also follow dataflow demand
             for (src, op) in h.all_linked_outputs(n) {
-                let needs_predecessor = match h.get_optype(src).port_kind(op).unwrap() {
-                    EdgeKind::Value(_) => {
-                        h.get_optype(src).is_load_constant()
-                            || results
-                                .try_read_wire_concrete::<Value, _, _>(Wire::new(src, op))
-                                .is_err()
-                    }
-                    EdgeKind::StateOrder | EdgeKind::Const(_) | EdgeKind::Function(_) => true,
-                    EdgeKind::ControlFlow => false, // we always include all children of a CFG above
-                    _ => true, // needed as EdgeKind non-exhaustive; not knowing what it is, assume the worst
-                };
-                if needs_predecessor {
-                    q.push_back(src);
+                if matches!(h.get_optype(src).port_kind(op).unwrap(), EdgeKind::Value(_))
+                    && results
+                        .try_read_wire_concrete::<Value, _, _>(Wire::new(src, op))
+                        .is_ok()
+                    && !h.get_optype(src).is_load_constant()
+                {
+                    continue;
                 }
+                // All other edge types --> we need the predecessors (even if included already e.g. ControlFlow)
+                q.push_back(src);
             }
         }
+        needed
     }
 }
 

hugr-passes/src/lib.rs

@@ -27,6 +27,8 @@ pub use monomorphize::monomorphize;
 pub use monomorphize::{MonomorphizeError, MonomorphizePass};
 pub mod nest_cfgs;
 pub mod non_local;
+mod static_eval;
+pub use static_eval::static_eval;
 pub mod validation;
 pub use force_order::{force_order, force_order_by_key};
 pub use lower::{lower_ops, replace_many_ops};