feat!: Extend LowerTypes pass to linearize by inserting copy/discard

hugr-passes/src/replace_types.rs

@@ -1,39 +1,47 @@
 #![allow(clippy::type_complexity)]
 #![warn(missing_docs)]
-//! Replace types with other types across the Hugr.
+//! Replace types with other types across the Hugr. See [ReplaceTypes] and [Linearizer].
 //!
-//! Parametrized types and ops will be reparametrized taking into account the replacements,
-//! but any ops taking/returning the replaced types *not* as a result of parametrization,
-//! will also need to be replaced - see [ReplaceTypes::replace_op]. (Similarly [Const]s.)
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::sync::Arc;
 
 use thiserror::Error;
 
+use hugr_core::builder::{BuildError, BuildHandle, Dataflow};
 use hugr_core::extension::{ExtensionId, OpDef, SignatureError, TypeDef};
 use hugr_core::hugr::hugrmut::HugrMut;
 use hugr_core::ops::constant::{OpaqueValue, Sum};
+use hugr_core::ops::handle::DataflowOpID;
 use hugr_core::ops::{
     AliasDefn, Call, CallIndirect, Case, Conditional, Const, DataflowBlock, ExitBlock, ExtensionOp,
-    FuncDecl, FuncDefn, Input, LoadConstant, LoadFunction, OpType, Output, Tag, TailLoop, Value,
-    CFG, DFG,
+    FuncDecl, FuncDefn, Input, LoadConstant, LoadFunction, OpTrait, OpType, Output, Tag, TailLoop,
+    Value, CFG, DFG,
 };
-use hugr_core::types::{CustomType, Transformable, Type, TypeArg, TypeEnum, TypeTransformer};
-use hugr_core::{Hugr, Node};
+use hugr_core::types::{
+    CustomType, Signature, Transformable, Type, TypeArg, TypeEnum, TypeTransformer,
+};
+use hugr_core::{Hugr, HugrView, Node, Wire};
 
 use crate::validation::{ValidatePassError, ValidationLevel};
 
-/// A thing with which an Op (i.e. node) can be replaced
+mod linearize;
+pub use linearize::{CallbackHandler, DelegatingLinearizer, LinearizeError, Linearizer};
+
+/// A recipe for creating a dataflow Node - as a new child of a [DataflowParent]
+/// or in order to replace an existing node.
+///
+/// [DataflowParent]: hugr_core::ops::OpTag::DataflowParent
 #[derive(Clone, Debug, PartialEq)]
-pub enum OpReplacement {
-    /// Keep the same node, change only the op (updating types of inputs/outputs)
+pub enum NodeTemplate {
+    /// A single node - so if replacing an existing node, change only the op
     SingleOp(OpType),
-    /// Defines a sub-Hugr to splice in place of the op - a [CFG], [Conditional], [DFG]
-    /// or [TailLoop], which must have the same inputs and outputs as the original op,
-    /// modulo replacement.
+    /// Defines a sub-Hugr to insert, whose root becomes (or replaces) the desired Node.
+    /// The root must be a [CFG], [Conditional], [DFG] or [TailLoop].
     // Not a FuncDefn, nor Case/DataflowBlock
-    /// Note this will be of limited use before [monomorphization](super::monomorphize()) because
-    /// the sub-Hugr will not be able to use type variables present in the op.
+    /// Note this will be of limited use before [monomorphization](super::monomorphize())
+    /// because the new subtree will not be able to use type variables present in the
+    /// parent Hugr or previous op.
     // TODO: store also a vec<TypeParam>, and update Hugr::validate to take &[TypeParam]s
     // (defaulting to empty list) - see https://github.com/CQCL/hugr/issues/709
     CompoundOp(Box<Hugr>),
@@ -42,12 +50,33 @@ pub enum OpReplacement {
     // So client should add the functions before replacement, then remove unused ones afterwards.)
 }
 
-impl OpReplacement {
+impl NodeTemplate {
+    /// Adds this instance to the specified [HugrMut] as a new node or subtree under a
+    /// given parent, returning the unique new child (of that parent) thus created
+    pub fn add_hugr(self, hugr: &mut impl HugrMut, parent: Node) -> Node {
+        match self {
+            NodeTemplate::SingleOp(op_type) => hugr.add_node_with_parent(parent, op_type),
+            NodeTemplate::CompoundOp(new_h) => hugr.insert_hugr(parent, *new_h).new_root,
+        }
+    }
+
+    /// Adds this instance to the specified [Dataflow] builder as a new node or subtree
+    pub fn add(
+        self,
+        dfb: &mut impl Dataflow,
+        inputs: impl IntoIterator<Item = Wire>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        match self {
+            NodeTemplate::SingleOp(opty) => dfb.add_dataflow_op(opty, inputs),
+            NodeTemplate::CompoundOp(h) => dfb.add_hugr_with_wires(*h, inputs),
+        }
+    }
+
     fn replace(&self, hugr: &mut impl HugrMut, n: Node) {
         assert_eq!(hugr.children(n).count(), 0);
         let new_optype = match self.clone() {
-            OpReplacement::SingleOp(op_type) => op_type,
-            OpReplacement::CompoundOp(new_h) => {
+            NodeTemplate::SingleOp(op_type) => op_type,
+            NodeTemplate::CompoundOp(new_h) => {
                 let new_root = hugr.insert_hugr(n, *new_h).new_root;
                 let children = hugr.children(new_root).collect::<Vec<_>>();
                 let root_opty = hugr.remove_node(new_root);
@@ -59,16 +88,50 @@ impl OpReplacement {
         };
         *hugr.optype_mut(n) = new_optype;
     }
+
+    fn signature(&self) -> Option<Cow<'_, Signature>> {
+        match self {
+            NodeTemplate::SingleOp(op_type) => op_type,
+            NodeTemplate::CompoundOp(hugr) => hugr.root_type(),
+        }
+        .dataflow_signature()
+    }
 }
 
 /// A configuration of what types, ops, and constants should be replaced with what.
 /// May be applied to a Hugr via [Self::run].
+///
+/// Parametrized types and ops will be reparametrized taking into account the
+/// replacements, but any ops taking/returning the replaced types *not* as a result of
+/// parametrization, will also need to be replaced - see [Self::replace_op].
+/// Similarly [Const]s.
+///
+/// Types that are [Copyable](hugr_core::types::TypeBound::Copyable) may also be replaced
+/// with types that are not, see [Linearizer].
+///
+/// Note that although this pass may be used before [monomorphization], there are some
+/// limitations (that do not apply if done after [monomorphization]):
+/// * [NodeTemplate::CompoundOp] only works for operations that do not use type variables
+/// * "Overrides" of specific instantiations of polymorphic types will not be detected if
+///   the instantiations are created inside polymorphic functions. For example, suppose
+///   we [Self::replace_type] type `A` with `X`, [Self::replace_parametrized_type]
+///   container `MyList` with `List`, and [Self::replace_type] `MyList<A>` with
+///   `SpecialListOfXs`. If a function `foo` polymorphic over a type variable `T` dealing
+///   with `MyList<T>`s, that is called with type argument `A`, then `foo<T>` will be
+///   updated to deal with `List<T>`s and the call `foo<A>` updated to `foo<X>`, but this
+///   will still result in using `List<X>` rather than `SpecialListOfXs`. (However this
+///   would be fine *after* [monomorphization]: the monomorphic definition of `foo_A`
+///   would use `SpecialListOfXs`.)
+/// * See also limitations noted for [Linearizer].
+///
+/// [monomorphization]: super::monomorphize()
 #[derive(Clone, Default)]
 pub struct ReplaceTypes {
     type_map: HashMap<CustomType, Type>,
     param_types: HashMap<ParametricType, Arc<dyn Fn(&[TypeArg]) -> Option<Type>>>,
-    op_map: HashMap<OpHashWrapper, OpReplacement>,
-    param_ops: HashMap<ParametricOp, Arc<dyn Fn(&[TypeArg]) -> Option<OpReplacement>>>,
+    linearize: DelegatingLinearizer,
+    op_map: HashMap<OpHashWrapper, NodeTemplate>,
+    param_ops: HashMap<ParametricOp, Arc<dyn Fn(&[TypeArg]) -> Option<NodeTemplate>>>,
     consts: HashMap<
         CustomType,
         Arc<dyn Fn(&OpaqueValue, &ReplaceTypes) -> Result<Value, ReplaceTypesError>>,
@@ -109,6 +172,8 @@ pub enum ReplaceTypesError {
     SignatureError(#[from] SignatureError),
     #[error(transparent)]
     ValidationError(#[from] ValidatePassError),
+    #[error(transparent)]
+    LinearizeError(#[from] LinearizeError),
 }
 
 impl ReplaceTypes {
@@ -157,16 +222,33 @@ impl ReplaceTypes {
         // (depending on arguments - i.e. if src's TypeDefBound is anything other than
         // `TypeDefBound::Explicit(TypeBound::Copyable)`) but that seems an annoying
         // overapproximation. Moreover, these depend upon the *return type* of the Fn.
+        // It would be too awkward to require:
+        // dest_fn: impl Fn(&TypeArg) -> (Type,
+        //                                Fn(&Linearizer) -> NodeTemplate, // copy
+        //                                Fn(&Linearizer) -> NodeTemplate)` // discard
         self.param_types.insert(src.into(), Arc::new(dest_fn));
     }
 
+    /// Allows to configure how to deal with types/wires that were [Copyable]
+    /// but have become linear as a result of type-changing. Specifically,
+    /// the [Linearizer] is used whenever lowering produces an outport which both
+    /// * has a non-[Copyable] type - perhaps a direct substitution, or perhaps e.g.
+    ///   as a result of changing the element type of a collection such as an [`array`]
+    /// * has other than one connected inport,
+    ///
+    /// [Copyable]: hugr_core::types::TypeBound::Copyable
+    /// [`array`]: hugr_core::std_extensions::collections::array::array_type
+    pub fn linearizer(&mut self) -> &mut DelegatingLinearizer {
+        &mut self.linearize
+    }
+
     /// Configures this instance to change occurrences of `src` to `dest`.
     /// Note that if `src` is an instance of a *parametrized* [OpDef], this takes
     /// precedence over [Self::replace_parametrized_op] where the `src`s overlap. Thus,
     /// this should only be used on already-*[monomorphize](super::monomorphize())d*
     /// Hugrs, as substitution (parametric polymorphism) happening later will not respect
     /// this replacement.
-    pub fn replace_op(&mut self, src: &ExtensionOp, dest: OpReplacement) {
+    pub fn replace_op(&mut self, src: &ExtensionOp, dest: NodeTemplate) {
         self.op_map.insert(OpHashWrapper::from(src), dest);
     }
 
@@ -179,7 +261,7 @@ impl ReplaceTypes {
     pub fn replace_parametrized_op(
         &mut self,
         src: &OpDef,
-        dest_fn: impl Fn(&[TypeArg]) -> Option<OpReplacement> + 'static,
+        dest_fn: impl Fn(&[TypeArg]) -> Option<NodeTemplate> + 'static,
     ) {
         self.param_ops.insert(src.into(), Arc::new(dest_fn));
     }
@@ -221,6 +303,22 @@ impl ReplaceTypes {
         let mut changed = false;
         for n in hugr.nodes().collect::<Vec<_>>() {
             changed |= self.change_node(hugr, n)?;
+            let new_dfsig = hugr.get_optype(n).dataflow_signature();
+            if let Some(new_sig) = new_dfsig
+                .filter(|_| changed && n != hugr.root())
+                .map(Cow::into_owned)
+            {
+                for outp in new_sig.output_ports() {
+                    if !new_sig.out_port_type(outp).unwrap().copyable() {
+                        let targets = hugr.linked_inputs(n, outp).collect::<Vec<_>>();
+                        if targets.len() != 1 {
+                            hugr.disconnect(n, outp);
+                            let src = Wire::new(n, outp);
+                            self.linearize.insert_copy_discard(hugr, src, &targets)?;
+                        }
+                    }
+                }
+            }
         }
         Ok(changed)
     }
@@ -452,7 +550,7 @@ mod test {
     use hugr_core::{hugr::IdentList, type_row, Extension, HugrView};
     use itertools::Itertools;
 
-    use super::{handlers::list_const, OpReplacement, ReplaceTypes};
+    use super::{handlers::list_const, NodeTemplate, ReplaceTypes};
 
     const PACKED_VEC: &str = "PackedVec";
     const READ: &str = "read";
@@ -513,7 +611,7 @@ mod test {
     }
 
     fn lowerer(ext: &Arc<Extension>) -> ReplaceTypes {
-        fn lowered_read(args: &[TypeArg]) -> Option<OpReplacement> {
+        fn lowered_read(args: &[TypeArg]) -> Option<NodeTemplate> {
             let ty = just_elem_type(args);
             let mut dfb = DFGBuilder::new(inout_sig(
                 vec![array_type(64, ty.clone()), i64_t()],
@@ -532,7 +630,7 @@ mod test {
             let [res] = dfb
                 .build_unwrap_sum(1, option_type(Type::from(ty.clone())), opt)
                 .unwrap();
-            Some(OpReplacement::CompoundOp(Box::new(
+            Some(NodeTemplate::CompoundOp(Box::new(
                 dfb.finish_hugr_with_outputs([res]).unwrap(),
             )))
         }
@@ -545,7 +643,7 @@ mod test {
         );
         lw.replace_op(
             &read_op(ext, bool_t()),
-            OpReplacement::SingleOp(
+            NodeTemplate::SingleOp(
                 ExtensionOp::new(ext.get_op("lowered_read_bool").unwrap().clone(), [])
                     .unwrap()
                     .into(),
@@ -824,7 +922,7 @@ mod test {
             e.get_op(READ).unwrap().as_ref(),
             Box::new(|args: &[TypeArg]| {
                 option_contents(just_elem_type(args)).map(|elem| {
-                    OpReplacement::SingleOp(
+                    NodeTemplate::SingleOp(
                         ListOp::get
                             .with_type(elem)
                             .to_extension_op()