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

hugr-passes/src/replace_types.rs

@@ -5,24 +5,30 @@
 //! 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::{Hugr, Node, Wire};
 
 use crate::validation::{ValidatePassError, ValidationLevel};
 
+mod linearize;
+pub use linearize::{LinearizeError, Linearizer};
+
 /// A thing with which an Op (i.e. node) can be replaced
 #[derive(Clone, Debug, PartialEq)]
 pub enum OpReplacement {
@@ -43,6 +49,27 @@ pub enum OpReplacement {
 }
 
 impl OpReplacement {
+    /// 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 {
+            OpReplacement::SingleOp(op_type) => hugr.add_node_with_parent(parent, op_type),
+            OpReplacement::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 {
+            OpReplacement::SingleOp(opty) => dfb.add_dataflow_op(opty, inputs),
+            OpReplacement::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() {
@@ -67,6 +94,7 @@ impl OpReplacement {
 pub struct ReplaceTypes {
     type_map: HashMap<CustomType, Type>,
     param_types: HashMap<ParametricType, Arc<dyn Fn(&[TypeArg]) -> Option<Type>>>,
+    linearize: Linearizer,
     op_map: HashMap<OpHashWrapper, OpReplacement>,
     param_ops: HashMap<ParametricOp, Arc<dyn Fn(&[TypeArg]) -> Option<OpReplacement>>>,
     consts: HashMap<
@@ -109,6 +137,8 @@ pub enum ReplaceTypesError {
     SignatureError(#[from] SignatureError),
     #[error(transparent)]
     ValidationError(#[from] ValidatePassError),
+    #[error(transparent)]
+    LinearizeError(#[from] LinearizeError),
 }
 
 impl ReplaceTypes {
@@ -157,9 +187,55 @@ 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) -> OpReplacement, // copy
+        //                                Fn(&Linearizer) -> OpReplacement)` // discard
         self.param_types.insert(src.into(), Arc::new(dest_fn));
     }
 
+    /// Configures this instance that, when an outport of type `src` has other than one connected
+    /// inport, the specified `copy` and or `discard` ops should be used to wire it to those inports.
+    /// (`copy` should have exactly one inport, of type `src`, and two outports, of same type;
+    /// `discard` should have exactly one inport, of type 'src', and no outports.)
+    ///
+    /// Also applies if `src` is an element of a Sum or other type.
+    ///
+    /// # Errors
+    ///
+    /// If `src` is [Copyable], it is returned as an `Err
+    ///
+    /// [Copyable]: hugr_core::types::TypeBound::Copyable
+    pub fn linearize(
+        &mut self,
+        src: Type,
+        copy: OpReplacement,
+        discard: OpReplacement,
+    ) -> Result<(), Type> {
+        self.linearize.register(src, copy, discard)
+    }
+
+    /// Configures this instance that when lowering produces an outport which
+    /// * has type an instantiation of the parametric type `src`, and
+    /// * is not [Copyable](hugr_core::types::TypeBound::Copyable), and
+    /// * has other than one connected inport,
+    ///
+    /// ...then these functions should be used to generate `copy` or `discard` ops.
+    ///
+    /// (That is, this is the equivalent of [Self::linearize] but for parametric types.)
+    ///
+    /// The [Linearizer] is passed so that the callbacks can use it to generate
+    /// `copy/`discard` ops for other types (e.g. the elements of a collection),
+    /// as part of an [OpReplacement::CompoundOp].
+    pub fn linearize_parametric(
+        &mut self,
+        src: &TypeDef,
+        copy_fn: Box<dyn Fn(&[TypeArg], &Linearizer) -> Result<OpReplacement, LinearizeError>>,
+        discard_fn: Box<dyn Fn(&[TypeArg], &Linearizer) -> Result<OpReplacement, LinearizeError>>,
+    ) {
+        self.linearize.register_parametric(src, copy_fn, discard_fn)
+    }
+
     /// 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,
@@ -221,6 +297,24 @@ 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) = (changed && n != hugr.root())
+                .then_some(new_dfsig)
+                .flatten()
+                .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 typ = new_sig.out_port_type(outp).unwrap();
+                            self.linearize
+                                .insert_copy_discard(hugr, n, outp, typ, &targets)?;
+                        }
+                    }
+                }
+            }
         }
         Ok(changed)
     }