feat: add verification to constant folding

hugr/src/algorithm.rs

@@ -4,3 +4,34 @@ pub mod const_fold;
 mod half_node;
 pub mod merge_bbs;
 pub mod nest_cfgs;
+
+#[derive(Debug, Clone, Copy, Ord, Eq, PartialOrd, PartialEq)]
+/// A type for algorithms to take as configuration, specifying how much
+/// verification they should do. Algorithms that accept this configuration
+/// should at least verify that input HUGRs are valid, and that output HUGRs are
+/// valid.
+///
+/// The default level is `None` because verification can be expensive.
+pub enum VerifyLevel {
+    /// Do no verification.
+    None,
+    /// Verify using [HugrView::validate_no_extensions]. This is useful when you
+    /// do not expect valid Extension annotations on Nodes.
+    ///
+    /// [HugrView::validate_no_extensions]: crate::HugrView::validate_no_extensions
+    WithoutExtensions,
+    /// Verify using [HugrView::validate].
+    ///
+    /// [HugrView::validate]: crate::HugrView::validate
+    WithExtensions,
+}
+
+impl Default for VerifyLevel {
+    fn default() -> Self {
+        if cfg!(test) {
+            Self::WithoutExtensions
+        } else {
+            Self::None
+        }
+    }
+}

hugr/src/algorithm/const_fold.rs

@@ -3,8 +3,11 @@
 use std::collections::{BTreeSet, HashMap};
 
 use itertools::Itertools;
+use thiserror::Error;
 
+use crate::hugr::{SimpleReplacementError, ValidationError};
 use crate::types::SumType;
+use crate::Direction;
 use crate::{
     builder::{DFGBuilder, Dataflow, DataflowHugr},
     extension::{ConstFoldResult, ExtensionRegistry},
@@ -19,6 +22,90 @@ use crate::{
     Hugr, HugrView, IncomingPort, Node, SimpleReplacement,
 };
 
+use super::VerifyLevel;
+
+#[derive(Error, Debug)]
+#[allow(missing_docs)]
+pub enum ConstFoldError {
+    #[error("Failed to verify {label} HUGR: {err}")]
+    VerifyError {
+        label: String,
+        #[source]
+        err: ValidationError,
+    },
+    #[error(transparent)]
+    SimpleReplaceError(#[from] SimpleReplacementError),
+}
+
+impl ConstFoldError {
+    fn verify_err(label: impl Into<String>, err: ValidationError) -> Self {
+        Self::VerifyError {
+            label: label.into(),
+            err,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+/// TODO
+pub struct ConstFoldConfig {
+    verify: VerifyLevel,
+}
+
+impl ConstFoldConfig {
+    /// Create a new `ConstFoldConfig` with default configuration.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Build a `ConstFoldConfig` with the given [VerifyLevel].
+    pub fn with_verify(mut self, verify: VerifyLevel) -> Self {
+        self.verify = verify;
+        self
+    }
+
+    fn verify_impl(
+        &self,
+        label: &str,
+        h: &impl HugrView,
+        reg: &ExtensionRegistry,
+    ) -> Result<(), ConstFoldError> {
+        match self.verify {
+            VerifyLevel::None => Ok(()),
+            VerifyLevel::WithoutExtensions => h.validate_no_extensions(reg),
+            VerifyLevel::WithExtensions => h.validate(reg),
+        }
+        .map_err(|err| ConstFoldError::verify_err(label, err))
+    }
+
+    /// Run the Constant Folding pass.
+    pub fn run(&self, h: &mut impl HugrMut, reg: &ExtensionRegistry) -> Result<(), ConstFoldError> {
+        self.verify_impl("input", h, reg)?;
+        loop {
+            // We can only safely apply a single replacement. Applying a
+            // replacement removes nodes and edges which may be referenced by
+            // further replacements returned by find_consts. Even worse, if we
+            // attempted to apply those replacements, expecting them to fail if
+            // the nodes and edges they reference had been deleted,  they may
+            // succeed because new nodes and edges reused the ids.
+            //
+            // We could be a lot smarter here, keeping track of `LoadConstant`
+            // nodes and only looking at their out neighbours.
+            let Some((replace, removes)) = find_consts(h, h.nodes(), reg).next() else {
+                break;
+            };
+            h.apply_rewrite(replace)?;
+            for rem in removes {
+                if let Ok(const_node) = h.apply_rewrite(rem) {
+                    // if the LoadConst was removed, try removing the Const too.
+                    let _ = h.apply_rewrite(RemoveConst(const_node));
+                }
+            }
+        }
+        self.verify_impl("output", h, reg)
+    }
+}
+
 /// Tag some output constants with [`OutgoingPort`] inferred from the ordering.
 fn out_row(consts: impl IntoIterator<Item = Value>) -> ConstFoldResult {
     let vec = consts
@@ -43,9 +130,10 @@ pub(crate) fn sorted_consts(consts: &[(IncomingPort, Value)]) -> Vec<&Value> {
         .map(|(_, c)| c)
         .collect()
 }
+
 /// For a given op and consts, attempt to evaluate the op.
 pub fn fold_leaf_op(op: &OpType, consts: &[(IncomingPort, Value)]) -> ConstFoldResult {
-    match op {
+    let fold_result = match op {
         OpType::Noop { .. } => out_row([consts.first()?.1.clone()]),
         OpType::MakeTuple { .. } => {
             out_row([Value::tuple(sorted_consts(consts).into_iter().cloned())])
@@ -69,7 +157,10 @@ pub fn fold_leaf_op(op: &OpType, consts: &[(IncomingPort, Value)]) -> ConstFoldR
             ext_op.constant_fold(consts)
         }
         _ => None,
-    }
+    };
+    assert!(fold_result.as_ref().map_or(true, |x| x.len()
+        == op.value_port_count(Direction::Outgoing)));
+    fold_result
 }
 
 /// Generate a graph that loads and outputs `consts` in order, validating
@@ -140,18 +231,16 @@ fn fold_op(
         })
         .unzip();
     // attempt to evaluate op
-    let folded = fold_leaf_op(neighbour_op, &in_consts)?;
-    let (op_outs, consts): (Vec<_>, Vec<_>) = folded.into_iter().unzip();
-    let nu_out = op_outs
+    let (nu_out, consts): (HashMap<_, _>, Vec<_>) = fold_leaf_op(neighbour_op, &in_consts)?
         .into_iter()
         .enumerate()
-        .filter_map(|(i, out)| {
-            // map from the ports the op was linked to, to the output ports of
-            // the replacement.
-            hugr.single_linked_input(op_node, out)
-                .map(|np| (np, i.into()))
+        .filter_map(|(i, (op_out, konst))| {
+            // for each used port of the op give the nu_out entry and the
+            // corresponding Value
+            hugr.single_linked_input(op_node, op_out)
+                .map(|np| ((np, i.into()), konst))
         })
-        .collect();
+        .unzip();
     let replacement = const_graph(consts, reg);
     let sibling_graph = SiblingSubgraph::try_from_nodes([op_node], hugr)
         .expect("Operation should form valid subgraph.");
@@ -172,39 +261,17 @@ fn get_const(hugr: &impl HugrView, op_node: Node, in_p: IncomingPort) -> Option<
     let (load_n, _) = hugr.single_linked_output(op_node, in_p)?;
     let load_op = hugr.get_optype(load_n).as_load_constant()?;
     let const_node = hugr
-        .linked_outputs(load_n, load_op.constant_port())
-        .exactly_one()
-        .ok()?
+        .single_linked_output(load_n, load_op.constant_port())?
         .0;
-
     let const_op = hugr.get_optype(const_node).as_const()?;
 
     // TODO avoid const clone here
     Some((const_op.as_ref().clone(), load_n))
 }
 
 /// Exhaustively apply constant folding to a HUGR.
-pub fn constant_fold_pass(h: &mut impl HugrMut, reg: &ExtensionRegistry) {
-    loop {
-        // would be preferable if the candidates were updated to be just the
-        // neighbouring nodes of those added.
-        let rewrites = find_consts(h, h.nodes(), reg).collect_vec();
-        if rewrites.is_empty() {
-            break;
-        }
-        for (replace, removes) in rewrites {
-            h.apply_rewrite(replace).unwrap();
-            for rem in removes {
-                if let Ok(const_node) = h.apply_rewrite(rem) {
-                    // if the LoadConst was removed, try removing the Const too.
-                    if h.apply_rewrite(RemoveConst(const_node)).is_err() {
-                        // const cannot be removed - no problem
-                        continue;
-                    }
-                }
-            }
-        }
-    }
+pub fn constant_fold_pass<H: HugrMut>(h: &mut H, reg: &ExtensionRegistry) {
+    ConstFoldConfig::default().run(h, reg).unwrap()
 }
 
 #[cfg(test)]
@@ -395,4 +462,44 @@ mod test {
         let expected = Value::false_val();
         assert_fully_folded(&h, &expected);
     }
+
+    #[test]
+    fn orphan_output() {
+        // pseudocode:
+        // x0 := bool(true)
+        // x1 := not(x0)
+        // x2 := or(x0,x1)
+        // output x2 == true;
+        //
+        // We arange things so that the `or` folds away first, leaving the not
+        // with no outputs.
+        use crate::hugr::NodeType;
+        use crate::ops::handle::NodeHandle;
+
+        let mut build = DFGBuilder::new(FunctionType::new(type_row![], vec![BOOL_T])).unwrap();
+        let true_wire = build.add_load_value(Value::true_val());
+        // this Not will be manually replaced
+        let orig_not = build.add_dataflow_op(NotOp, [true_wire]).unwrap();
+        let r = build
+            .add_dataflow_op(
+                NaryLogic::Or.with_n_inputs(2),
+                [true_wire, orig_not.out_wire(0)],
+            )
+            .unwrap();
+        let or_node = r.node();
+        let parent = build.dfg_node;
+        let reg =
+            ExtensionRegistry::try_new([PRELUDE.to_owned(), logic::EXTENSION.to_owned()]).unwrap();
+        let mut h = build.finish_hugr_with_outputs(r.outputs(), &reg).unwrap();
+
+        // we delete the original Not and create a new One. This means it will be
+        // traversed by `constant_fold_pass` after the Or.
+        let new_not = h.add_node_with_parent(parent, NodeType::new_auto(NotOp));
+        h.connect(true_wire.node(), true_wire.source(), new_not, 0);
+        h.disconnect(or_node, IncomingPort::from(1));
+        h.connect(new_not, 0, or_node, 1);
+        h.remove_node(orig_not.node());
+        constant_fold_pass(&mut h, &reg);
+        assert_fully_folded(&h, &Value::true_val())
+    }
 }

hugr/src/hugr/views.rs

@@ -24,7 +24,10 @@ use itertools::{Itertools, MapInto};
 use portgraph::render::{DotFormat, MermaidFormat};
 use portgraph::{multiportgraph, LinkView, MultiPortGraph, PortView};
 
-use super::{Hugr, HugrError, NodeMetadata, NodeMetadataMap, NodeType, DEFAULT_NODETYPE};
+use super::{
+    Hugr, HugrError, NodeMetadata, NodeMetadataMap, NodeType, ValidationError, DEFAULT_NODETYPE,
+};
+use crate::extension::ExtensionRegistry;
 use crate::ops::handle::NodeHandle;
 use crate::ops::{OpParent, OpTag, OpTrait, OpType};
 
@@ -460,6 +463,18 @@ pub trait HugrView: sealed::HugrInternals {
         self.value_types(node, Direction::Outgoing)
             .map(|(p, t)| (p.as_outgoing().unwrap(), t))
     }
+
+    /// Check the validity of the underlying HUGR.
+    fn validate(&self, reg: &ExtensionRegistry) -> Result<(), ValidationError> {
+        self.base_hugr().validate(reg)
+    }
+
+    /// Check the validity of the underlying HUGR, but don't check consistency
+    /// of extension requirements between connected nodes or between parents and
+    /// children.
+    fn validate_no_extensions(&self, reg: &ExtensionRegistry) -> Result<(), ValidationError> {
+        self.base_hugr().validate_no_extensions(reg)
+    }
 }
 
 /// Wraps an iterator over [Port]s that are known to be [OutgoingPort]s

hugr/src/std_extensions/arithmetic/int_ops/const_fold.rs

@@ -16,6 +16,16 @@ use crate::{
 
 use super::IntOpDef;
 
+use lazy_static::lazy_static;
+
+lazy_static! {
+    static ref INARROW_ERROR_VALUE: Value = ConstError {
+        signal: 0,
+        message: "Integer too large to narrow".to_string(),
+    }
+    .into();
+}
+
 fn bitmask_from_width(width: u64) -> u64 {
     debug_assert!(width <= 64);
     if width == 64 {
@@ -111,28 +121,22 @@ pub(super) fn set_fold(op: &IntOpDef, def: &mut OpDef) {
                     let logwidth0: u8 = get_log_width(arg0).ok()?;
                     let logwidth1: u8 = get_log_width(arg1).ok()?;
                     let n0: &ConstInt = get_single_input_value(consts)?;
+                    (logwidth0 >= logwidth1 && n0.log_width() == logwidth0).then_some(())?;
 
                     let int_out_type = INT_TYPES[logwidth1 as usize].to_owned();
                     let sum_type = sum_with_error(int_out_type.clone());
-                    let err_value = || {
-                        let err_val = ConstError {
-                            signal: 0,
-                            message: "Integer too large to narrow".to_string(),
-                        };
-                        Value::sum(1, [err_val.into()], sum_type.clone())
+
+                    let mk_out_const = |i, mb_v: Result<Value, _>| {
+                        mb_v.and_then(|v| Value::sum(i, [v], sum_type))
                             .unwrap_or_else(|e| panic!("Invalid computed sum, {}", e))
                     };
                     let n0val: u64 = n0.value_u();
                     let out_const: Value = if n0val >> (1 << logwidth1) != 0 {
-                        err_value()
+                        mk_out_const(1, Ok(INARROW_ERROR_VALUE.clone()))
                     } else {
-                        Value::extension(ConstInt::new_u(logwidth1, n0val).unwrap())
+                        mk_out_const(0, ConstInt::new_u(logwidth1, n0val).map(Into::into))
                     };
-                    if logwidth0 < logwidth1 || n0.log_width() != logwidth0 {
-                        None
-                    } else {
-                        Some(vec![(0.into(), out_const)])
-                    }
+                    Some(vec![(0.into(), out_const)])
                 },
             ),
         },
@@ -145,29 +149,22 @@ pub(super) fn set_fold(op: &IntOpDef, def: &mut OpDef) {
                     let logwidth0: u8 = get_log_width(arg0).ok()?;
                     let logwidth1: u8 = get_log_width(arg1).ok()?;
                     let n0: &ConstInt = get_single_input_value(consts)?;
+                    (logwidth0 >= logwidth1 && n0.log_width() == logwidth0).then_some(())?;
 
                     let int_out_type = INT_TYPES[logwidth1 as usize].to_owned();
                     let sum_type = sum_with_error(int_out_type.clone());
-                    let err_value = || {
-                        let err_val = ConstError {
-                            signal: 0,
-                            message: "Integer too large to narrow".to_string(),
-                        };
-                        Value::sum(1, [err_val.into()], sum_type.clone())
+                    let mk_out_const = |i, mb_v: Result<Value, _>| {
+                        mb_v.and_then(|v| Value::sum(i, [v], sum_type))
                             .unwrap_or_else(|e| panic!("Invalid computed sum, {}", e))
                     };
                     let n0val: i64 = n0.value_s();
                     let ub = 1i64 << ((1 << logwidth1) - 1);
                     let out_const: Value = if n0val >= ub || n0val < -ub {
-                        err_value()
+                        mk_out_const(1, Ok(INARROW_ERROR_VALUE.clone()))
                     } else {
-                        Value::extension(ConstInt::new_s(logwidth1, n0val).unwrap())
+                        mk_out_const(0, ConstInt::new_s(logwidth1, n0val).map(Into::into))
                     };
-                    if logwidth0 < logwidth1 || n0.log_width() != logwidth0 {
-                        None
-                    } else {
-                        Some(vec![(0.into(), out_const)])
-                    }
+                    Some(vec![(0.into(), out_const)])
                 },
             ),
         },