feat: Constant-folding of integer and logic operations

hugr/src/extension.rs

@@ -37,7 +37,7 @@ mod const_fold;
 pub mod prelude;
 pub mod simple_op;
 pub mod validate;
-pub use const_fold::{ConstFold, ConstFoldResult};
+pub use const_fold::{ConstFold, ConstFoldResult, Folder};
 pub use prelude::{PRELUDE, PRELUDE_REGISTRY};
 
 pub mod declarative;

hugr/src/extension/const_fold.rs

@@ -2,8 +2,10 @@ use std::fmt::Formatter;
 
 use std::fmt::Debug;
 
+use crate::ops::Value;
 use crate::types::TypeArg;
 
+use crate::IncomingPort;
 use crate::OutgoingPort;
 
 use crate::ops;
@@ -45,3 +47,17 @@ where
         self(consts)
     }
 }
+
+type FoldFn = dyn Fn(&[TypeArg], &[(IncomingPort, Value)]) -> ConstFoldResult + Send + Sync;
+
+/// Type holding a boxed const-folding function.
+pub struct Folder {
+    /// Const-folding function.
+    pub folder: Box<FoldFn>,
+}
+
+impl ConstFold for Folder {
+    fn fold(&self, type_args: &[TypeArg], consts: &[(IncomingPort, Value)]) -> ConstFoldResult {
+        (self.folder)(type_args, consts)
+    }
+}

hugr/src/hugr/rewrite/inline_dfg.rs

@@ -206,12 +206,12 @@ mod test {
             )?;
             let [a] = inner.input_wires_arr();
             let c1 = c1.unwrap_or_else(|| make_const(&mut inner))?;
-            let a1 = inner.add_dataflow_op(IntOpDef::iadd.with_width(6), [a, c1])?;
+            let a1 = inner.add_dataflow_op(IntOpDef::iadd.with_log_width(6), [a, c1])?;
             inner.finish_with_outputs(a1.outputs())?
         };
         let [a1] = inner.outputs_arr();
 
-        let a1_sub_b = outer.add_dataflow_op(IntOpDef::isub.with_width(6), [a1, b])?;
+        let a1_sub_b = outer.add_dataflow_op(IntOpDef::isub.with_log_width(6), [a1, b])?;
         let mut outer = outer.finish_hugr_with_outputs(a1_sub_b.outputs(), &reg)?;
 
         // Sanity checks

hugr/src/ops/constant.rs

@@ -12,7 +12,10 @@ use itertools::Itertools;
 use smol_str::SmolStr;
 use thiserror::Error;
 
-pub use custom::{downcast_equal_consts, CustomConst, CustomSerialized};
+pub use custom::{
+    downcast_equal_consts, get_pair_of_input_values, get_single_input_value, CustomConst,
+    CustomSerialized,
+};
 
 #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
 /// An operation returning a constant value.

hugr/src/ops/constant/custom.rs

@@ -12,6 +12,9 @@ use crate::extension::ExtensionSet;
 use crate::macros::impl_box_clone;
 
 use crate::types::{CustomCheckFailure, Type};
+use crate::IncomingPort;
+
+use super::Value;
 
 use super::ValueName;
 
@@ -118,3 +121,21 @@ impl PartialEq for dyn CustomConst {
         (*self).equal_consts(other)
     }
 }
+
+/// Given a singleton list of constant operations, return the value.
+pub fn get_single_input_value<T: CustomConst>(consts: &[(IncomingPort, Value)]) -> Option<&T> {
+    let [(_, c)] = consts else {
+        return None;
+    };
+    c.get_custom_value()
+}
+
+/// Given a list of two constant operations, return the values.
+pub fn get_pair_of_input_values<T: CustomConst>(
+    consts: &[(IncomingPort, Value)],
+) -> Option<(&T, &T)> {
+    let [(_, c0), (_, c1)] = consts else {
+        return None;
+    };
+    Some((c0.get_custom_value()?, c1.get_custom_value()?))
+}

hugr/src/std_extensions/arithmetic/conversions.rs

@@ -76,18 +76,18 @@ impl MakeOpDef for ConvertOpDef {
 
 impl ConvertOpDef {
     /// Initialise a conversion op with an integer log width type argument.
-    pub fn with_width(self, log_width: u8) -> ConvertOpType {
+    pub fn with_log_width(self, log_width: u8) -> ConvertOpType {
         ConvertOpType {
             def: self,
-            log_width: log_width as u64,
+            log_width,
         }
     }
 }
-/// Concrete convert operation with integer width set.
+/// Concrete convert operation with integer log width set.
 #[derive(Debug, Clone, PartialEq)]
 pub struct ConvertOpType {
     def: ConvertOpDef,
-    log_width: u64,
+    log_width: u8,
 }
 
 impl NamedOp for ConvertOpType {
@@ -99,18 +99,20 @@ impl NamedOp for ConvertOpType {
 impl MakeExtensionOp for ConvertOpType {
     fn from_extension_op(ext_op: &ExtensionOp) -> Result<Self, OpLoadError> {
         let def = ConvertOpDef::from_def(ext_op.def())?;
-        let width = match *ext_op.args() {
+        let log_width: u64 = match *ext_op.args() {
             [TypeArg::BoundedNat { n }] => n,
             _ => return Err(SignatureError::InvalidTypeArgs.into()),
         };
         Ok(Self {
             def,
-            log_width: width,
+            log_width: u8::try_from(log_width).unwrap(),
         })
     }
 
     fn type_args(&self) -> Vec<crate::types::TypeArg> {
-        vec![TypeArg::BoundedNat { n: self.log_width }]
+        vec![TypeArg::BoundedNat {
+            n: self.log_width as u64,
+        }]
     }
 }
 

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

@@ -1,14 +1,15 @@
+use crate::ops::constant::get_single_input_value;
 use crate::ops::Value;
+use crate::std_extensions::arithmetic::int_types::INT_TYPES;
 use crate::{
     extension::{
         prelude::{sum_with_error, ConstError},
         ConstFold, ConstFoldResult, OpDef,
     },
     ops,
-    ops::constant::CustomConst,
     std_extensions::arithmetic::{
         float_types::ConstF64,
-        int_types::{get_log_width, ConstInt, INT_TYPES},
+        int_types::{get_log_width, ConstInt},
     },
     types::ConstTypeError,
     IncomingPort,
@@ -27,19 +28,12 @@ pub(super) fn set_fold(op: &ConvertOpDef, def: &mut OpDef) {
     }
 }
 
-fn get_input<T: CustomConst>(consts: &[(IncomingPort, ops::Value)]) -> Option<&T> {
-    let [(_, c)] = consts else {
-        return None;
-    };
-    c.get_custom_value()
-}
-
 fn fold_trunc(
     type_args: &[crate::types::TypeArg],
     consts: &[(IncomingPort, Value)],
     convert: impl Fn(f64, u8) -> Result<Value, ConstTypeError>,
 ) -> ConstFoldResult {
-    let f: &ConstF64 = get_input(consts)?;
+    let f: &ConstF64 = get_single_input_value(consts)?;
     let f = f.value();
     let [arg] = type_args else {
         return None;
@@ -105,7 +99,7 @@ impl ConstFold for ConvertU {
         _type_args: &[crate::types::TypeArg],
         consts: &[(IncomingPort, ops::Value)],
     ) -> ConstFoldResult {
-        let u: &ConstInt = get_input(consts)?;
+        let u: &ConstInt = crate::ops::constant::get_single_input_value(consts)?;
         let f = u.value_u() as f64;
         Some(vec![(0.into(), ConstF64::new(f).into())])
     }
@@ -119,7 +113,7 @@ impl ConstFold for ConvertS {
         _type_args: &[crate::types::TypeArg],
         consts: &[(IncomingPort, ops::Value)],
     ) -> ConstFoldResult {
-        let u: &ConstInt = get_input(consts)?;
+        let u: &ConstInt = get_single_input_value(consts)?;
         let f = u.value_s() as f64;
         Some(vec![(0.into(), ConstF64::new(f).into())])
     }