feat: ReplaceTypes: handlers for array constants + linearization

hugr-core/src/std_extensions/collections/array.rs

@@ -178,7 +178,9 @@ lazy_static! {
     };
 }
 
-fn array_type_def() -> &'static TypeDef {
+/// Gets the [TypeDef] for arrays. Note that instantiations are more easily
+/// created via [array_type] and [array_type_parametric]
+pub fn array_type_def() -> &'static TypeDef {
     EXTENSION.get_type(&ARRAY_TYPENAME).unwrap()
 }
 

hugr-passes/src/replace_types.rs

@@ -19,7 +19,7 @@ use hugr_core::ops::{
     Value, CFG, DFG,
 };
 use hugr_core::types::{
-    CustomType, Signature, Transformable, Type, TypeArg, TypeEnum, TypeTransformer,
+    ConstTypeError, CustomType, Signature, Transformable, Type, TypeArg, TypeEnum, TypeTransformer,
 };
 use hugr_core::{Hugr, HugrView, Node, Wire};
 
@@ -173,6 +173,8 @@ pub enum ReplaceTypesError {
     #[error(transparent)]
     ValidationError(#[from] ValidatePassError),
     #[error(transparent)]
+    ConstError(#[from] ConstTypeError),
+    #[error(transparent)]
     LinearizeError(#[from] LinearizeError),
 }
 
@@ -447,38 +449,7 @@ impl ReplaceTypes {
     }
 }
 
-pub mod handlers {
-    //! Callbacks for use with [ReplaceTypes::replace_consts_parametrized]
-    use hugr_core::ops::{constant::OpaqueValue, Value};
-    use hugr_core::std_extensions::collections::list::ListValue;
-    use hugr_core::types::Transformable;
-
-    use super::{ReplaceTypes, ReplaceTypesError};
-
-    /// Handler for [ListValue] constants that recursively [ReplaceTypes::change_value]s
-    /// the elements of the list
-    pub fn list_const(
-        val: &OpaqueValue,
-        repl: &ReplaceTypes,
-    ) -> Result<Option<Value>, ReplaceTypesError> {
-        let Some(lv) = val.value().downcast_ref::<ListValue>() else {
-            return Ok(None);
-        };
-        let mut vals: Vec<Value> = lv.get_contents().to_vec();
-        let mut ch = false;
-        for v in vals.iter_mut() {
-            ch |= repl.change_value(v)?;
-        }
-        // If none of the values has changed, assume the Type hasn't (Values have a single known type)
-        if !ch {
-            return Ok(None);
-        };
-
-        let mut elem_t = lv.get_element_type().clone();
-        elem_t.transform(repl)?;
-        Ok(Some(ListValue::new(elem_t, vals).into()))
-    }
-}
+pub mod handlers;
 
 #[derive(Clone, Hash, PartialEq, Eq)]
 struct OpHashWrapper {
@@ -536,20 +507,26 @@ mod test {
     use hugr_core::extension::simple_op::MakeExtensionOp;
     use hugr_core::extension::{TypeDefBound, Version};
 
+    use hugr_core::ops::constant::OpaqueValue;
     use hugr_core::ops::{ExtensionOp, NamedOp, OpTrait, OpType, Tag, Value};
     use hugr_core::std_extensions::arithmetic::int_types::ConstInt;
     use hugr_core::std_extensions::arithmetic::{conversions::ConvertOpDef, int_types::INT_TYPES};
     use hugr_core::std_extensions::collections::array::{
-        array_type, ArrayOp, ArrayOpDef, ArrayValue,
+        array_type, array_type_def, ArrayOp, ArrayOpDef, ArrayValue,
     };
     use hugr_core::std_extensions::collections::list::{
         list_type, list_type_def, ListOp, ListValue,
     };
 
+    use hugr_core::hugr::ValidationError;
     use hugr_core::types::{PolyFuncType, Signature, SumType, Type, TypeArg, TypeBound, TypeRow};
     use hugr_core::{hugr::IdentList, type_row, Extension, HugrView};
     use itertools::Itertools;
+    use rstest::rstest;
 
+    use crate::validation::ValidatePassError;
+
+    use super::ReplaceTypesError;
     use super::{handlers::list_const, NodeTemplate, ReplaceTypes};
 
     const PACKED_VEC: &str = "PackedVec";
@@ -951,4 +928,38 @@ mod test {
             ["NoBoundsCheck.read", "collections.list.get"]
         );
     }
+
+    #[rstest]
+    #[case(&[])]
+    #[case(&[3])]
+    #[case(&[5,7,11,13,17,19])]
+    fn array_const(#[case] vals: &[u64]) {
+        use super::handlers::array_const;
+        let mut dfb = DFGBuilder::new(inout_sig(
+            type_row![],
+            array_type(vals.len() as _, usize_t()),
+        ))
+        .unwrap();
+        let c = dfb.add_load_value(ArrayValue::new(
+            usize_t(),
+            vals.iter().map(|u| ConstUsize::new(*u).into()),
+        ));
+        let backup = dfb.finish_hugr_with_outputs([c]).unwrap();
+
+        let mut repl = ReplaceTypes::default();
+        let usize_custom_t = usize_t().as_extension().unwrap().clone();
+        repl.replace_type(usize_custom_t.clone(), INT_TYPES[6].clone());
+        repl.replace_consts(usize_custom_t, |cst: &OpaqueValue, _| {
+            let cu = cst.value().downcast_ref::<ConstUsize>().unwrap();
+            Ok(ConstInt::new_u(6, cu.value())?.into())
+        });
+        assert!(
+            matches!(repl.run(&mut backup.clone()), Err(ReplaceTypesError::ValidationError(ValidatePassError::OutputError {
+                err: ValidationError::IncompatiblePorts {from, to, ..}, ..
+            })) if backup.get_optype(from).is_const() && to == c.node())
+        );
+        repl.replace_consts_parametrized(array_type_def(), array_const);
+        let mut h = backup;
+        repl.run(&mut h).unwrap(); // Includes validation
+    }
 }

hugr-passes/src/replace_types/handlers.rs

@@ -0,0 +1,262 @@
+//! Callbacks for use with [ReplaceTypes::replace_consts_parametrized]
+//! and [DelegatingLinearizer::register_callback](super::DelegatingLinearizer::register_callback)
+
+use hugr_core::builder::{endo_sig, inout_sig, DFGBuilder, Dataflow, DataflowHugr};
+use hugr_core::extension::prelude::{option_type, UnwrapBuilder};
+use hugr_core::extension::ExtensionSet;
+use hugr_core::ops::{constant::OpaqueValue, Value};
+use hugr_core::ops::{OpTrait, OpType, Tag};
+use hugr_core::std_extensions::arithmetic::conversions::ConvertOpDef;
+use hugr_core::std_extensions::arithmetic::int_ops::IntOpDef;
+use hugr_core::std_extensions::arithmetic::int_types::{ConstInt, INT_TYPES};
+use hugr_core::std_extensions::collections::array::{
+    array_type, ArrayOpDef, ArrayRepeat, ArrayScan, ArrayValue,
+};
+use hugr_core::std_extensions::collections::list::ListValue;
+use hugr_core::types::{SumType, Transformable, Type, TypeArg};
+use hugr_core::{type_row, Hugr, HugrView};
+use itertools::Itertools;
+
+use super::{
+    CallbackHandler, LinearizeError, Linearizer, NodeTemplate, ReplaceTypes, ReplaceTypesError,
+};
+
+/// Handler for [ListValue] constants that updates the element type and
+/// recursively [ReplaceTypes::change_value]s the elements of the list
+pub fn list_const(
+    val: &OpaqueValue,
+    repl: &ReplaceTypes,
+) -> Result<Option<Value>, ReplaceTypesError> {
+    let Some(lv) = val.value().downcast_ref::<ListValue>() else {
+        return Ok(None);
+    };
+    let mut elem_t = lv.get_element_type().clone();
+    if !elem_t.transform(repl)? {
+        // No change to type, so values should not change either
+        return Ok(None);
+    }
+
+    let mut vals: Vec<Value> = lv.get_contents().to_vec();
+    for v in vals.iter_mut() {
+        repl.change_value(v)?;
+    }
+    Ok(Some(ListValue::new(elem_t, vals).into()))
+}
+
+/// Handler for [ArrayValue] constants that recursively
+/// [ReplaceTypes::change_value]s the elements of the list
+pub fn array_const(
+    val: &OpaqueValue,
+    repl: &ReplaceTypes,
+) -> Result<Option<Value>, ReplaceTypesError> {
+    let Some(av) = val.value().downcast_ref::<ArrayValue>() else {
+        return Ok(None);
+    };
+    let mut elem_t = av.get_element_type().clone();
+    if !elem_t.transform(repl)? {
+        // No change to type, so values should not change either
+        return Ok(None);
+    }
+
+    let mut vals: Vec<Value> = av.get_contents().to_vec();
+    for v in vals.iter_mut() {
+        repl.change_value(v)?;
+    }
+    Ok(Some(ArrayValue::new(elem_t, vals).into()))
+}
+
+fn runtime_reqs(h: &Hugr) -> ExtensionSet {
+    h.signature(h.root()).unwrap().runtime_reqs.clone()
+}
+
+/// Handler for copying/discarding arrays, for use with [register_callback] for
+/// [array_type_def](hugr_core::std_extensions::collections::array::array_type_def)
+///
+/// [register_callback]: super::DelegatingLinearizer::register_callback
+pub fn linearize_array(
+    args: &[TypeArg],
+    num_outports: usize,
+    lin: &CallbackHandler,
+) -> Result<NodeTemplate, LinearizeError> {
+    // Require known length i.e. usable only after monomorphization, due to no-variables limitation
+    // restriction on NodeTemplate::CompoundOp
+    let [TypeArg::BoundedNat { n }, TypeArg::Type { ty }] = args else {
+        panic!("Illegal TypeArgs to array: {:?}", args)
+    };
+    if num_outports == 0 {
+        // "Simple" discard - first map each element to unit (via type-specific discard):
+        let map_fn = {
+            let mut dfb = DFGBuilder::new(inout_sig(ty.clone(), Type::UNIT)).unwrap();
+            let [to_discard] = dfb.input_wires_arr();
+            lin.copy_discard_op(ty, 0)?
+                .add(&mut dfb, [to_discard])
+                .unwrap();
+            let ret = dfb.add_load_value(Value::unary_unit_sum());
+            dfb.finish_hugr_with_outputs([ret]).unwrap()
+        };
+        // Now array.scan that over the input array to get an array of unit (which can be discarded)
+        let array_scan = ArrayScan::new(ty.clone(), Type::UNIT, vec![], *n, runtime_reqs(&map_fn));
+        let in_type = array_type(*n, ty.clone());
+        return Ok(NodeTemplate::CompoundOp(Box::new({
+            let mut dfb = DFGBuilder::new(inout_sig(in_type, type_row![])).unwrap();
+            let [in_array] = dfb.input_wires_arr();
+            let map_fn = dfb.add_load_value(Value::Function {
+                hugr: Box::new(map_fn),
+            });
+            // scan has one output, an array of unit, so just ignore/discard that
+            dfb.add_dataflow_op(array_scan, [in_array, map_fn]).unwrap();
+            dfb.finish_hugr_with_outputs([]).unwrap()
+        })));
+    };
+    // The num_outports>1 case will simplify, and unify with the previous, when we have a
+    // more general ArrayScan https://github.com/CQCL/hugr/issues/2041. In the meantime:
+    let num_new = num_outports - 1;
+    let array_ty = array_type(*n, ty.clone());
+    let mut dfb = DFGBuilder::new(inout_sig(
+        array_ty.clone(),
+        vec![array_ty.clone(); num_outports],
+    ))
+    .unwrap();
+    // 1. make num_new array<SZ, Option<T>>, initialized to None...
+    let option_sty = option_type(ty.clone());
+    let option_ty = Type::from(option_sty.clone());
+    let arrays_of_none = {
+        let fn_none = {
+            let mut dfb = DFGBuilder::new(inout_sig(vec![], option_ty.clone())).unwrap();
+            let none = dfb
+                .add_dataflow_op(Tag::new(0, vec![type_row![], ty.clone().into()]), [])
+                .unwrap();
+            dfb.finish_hugr_with_outputs(none.outputs()).unwrap()
+        };
+        let repeats =
+            vec![ArrayRepeat::new(option_ty.clone(), *n, runtime_reqs(&fn_none)); num_new];
+        let fn_none = dfb.add_load_value(Value::function(fn_none).unwrap());
+        repeats
+            .into_iter()
+            .map(|rpt| {
+                let [arr] = dfb.add_dataflow_op(rpt, [fn_none]).unwrap().outputs_arr();
+                arr
+            })
+            .collect::<Vec<_>>()
+    };
+
+    // 2. use a scan through the input array, copying the element num_outputs times;
+    // return the first copy, and put each of the other copies into one of the array<option>
+
+    let i64_t = INT_TYPES[6].to_owned();
+    let option_array = array_type(*n, option_ty.clone());
+    let copy_elem = {
+        let mut io = vec![ty.clone(), i64_t.clone()];
+        io.extend(vec![option_array.clone(); num_new]);
+        let mut dfb = DFGBuilder::new(endo_sig(io)).unwrap();
+        let mut inputs = dfb.input_wires();
+        let elem = inputs.next().unwrap();
+        let idx = inputs.next().unwrap();
+        let opt_arrays = inputs.collect::<Vec<_>>();
+        let [idx_usz] = dfb
+            .add_dataflow_op(ConvertOpDef::itousize.without_log_width(), [idx])
+            .unwrap()
+            .outputs_arr();
+        let mut copies = lin
+            .copy_discard_op(ty, num_outports)?
+            .add(&mut dfb, [elem])
+            .unwrap()
+            .outputs();
+        let copy0 = copies.next().unwrap(); // We'll return this directly
+                                            // Wrap remaining copies into an option
+        let set_op = OpType::from(ArrayOpDef::set.to_concrete(option_ty.clone(), *n));
+        let either_st = set_op.dataflow_signature().unwrap().output[0]
+            .as_sum()
+            .unwrap()
+            .clone();
+        let opt_arrays = opt_arrays
+            .into_iter()
+            .zip_eq(copies)
+            .map(|(opt_array, copy1)| {
+                let [tag] = dfb
+                    .add_dataflow_op(Tag::new(1, vec![type_row![], ty.clone().into()]), [copy1])
+                    .unwrap()
+                    .outputs_arr();
+                let [set_result] = dfb
+                    .add_dataflow_op(set_op.clone(), [opt_array, idx_usz, tag])
+                    .unwrap()
+                    .outputs_arr();
+                // set should always be successful
+                let [none, opt_array] = dfb
+                    .build_unwrap_sum(1, either_st.clone(), set_result)
+                    .unwrap();
+                //the removed element is an option, which should always be none (and thus discardable)
+                let [] = dfb
+                    .build_unwrap_sum(0, SumType::new_option(ty.clone()), none)
+                    .unwrap();
+                opt_array
+            })
+            .collect::<Vec<_>>(); // stop borrowing dfb
+
+        let cst1 = dfb.add_load_value(ConstInt::new_u(6, 1).unwrap());
+        let [new_idx] = dfb
+            .add_dataflow_op(IntOpDef::iadd.with_log_width(6), [idx, cst1])
+            .unwrap()
+            .outputs_arr();
+        dfb.finish_hugr_with_outputs([copy0, new_idx].into_iter().chain(opt_arrays))
+            .unwrap()
+    };
+    let [in_array] = dfb.input_wires_arr();
+    let scan1 = ArrayScan::new(
+        ty.clone(),
+        ty.clone(),
+        std::iter::once(i64_t)
+            .chain(vec![option_array; num_new])
+            .collect(),
+        *n,
+        runtime_reqs(&copy_elem),
+    );
+
+    let copy_elem = dfb.add_load_value(Value::function(copy_elem).unwrap());
+    let cst0 = dfb.add_load_value(ConstInt::new_u(6, 0).unwrap());
+
+    let mut outs = dfb
+        .add_dataflow_op(
+            scan1,
+            [in_array, copy_elem, cst0]
+                .into_iter()
+                .chain(arrays_of_none),
+        )
+        .unwrap()
+        .outputs();
+    let out_array1 = outs.next().unwrap();
+    let _idx_out = outs.next().unwrap();
+    let opt_arrays = outs;
+
+    //3. Scan each array-of-options, 'unwrapping' each element into a non-option
+    let unwrap_elem = {
+        let mut dfb =
+            DFGBuilder::new(inout_sig(Type::from(option_ty.clone()), ty.clone())).unwrap();
+        let [opt] = dfb.input_wires_arr();
+        let [val] = dfb.build_unwrap_sum(1, option_sty.clone(), opt).unwrap();
+        dfb.finish_hugr_with_outputs([val]).unwrap()
+    };
+
+    let unwrap_scan = ArrayScan::new(
+        option_ty.clone(),
+        ty.clone(),
+        vec![],
+        *n,
+        runtime_reqs(&unwrap_elem),
+    );
+    let unwrap_elem = dfb.add_load_value(Value::function(unwrap_elem).unwrap());
+
+    let out_arrays = std::iter::once(out_array1)
+        .chain(opt_arrays.map(|opt_array| {
+            let [out_array] = dfb
+                .add_dataflow_op(unwrap_scan.clone(), [opt_array, unwrap_elem])
+                .unwrap()
+                .outputs_arr();
+            out_array
+        }))
+        .collect::<Vec<_>>();
+
+    Ok(NodeTemplate::CompoundOp(Box::new(
+        dfb.finish_hugr_with_outputs(out_arrays).unwrap(),
+    )))
+}