feat: Add a borrow_array type replacement pass (#975)

Draft because performance benchmarks with Guppy using this are needed
first to see if this is suitable. There are probably a lot of ways that
some of the more repetitive code in all the `dest` functions could be
generalised but for now I prioritised getting something working.

Closes [#2397](Quantinuum/hugr#2397) (but should add
a new issue for doing this properly in LLVM at some point)

Author: tatiana-s

tket-qsystem/src/lib.rs

@@ -9,6 +9,7 @@ pub mod llvm;
 mod lower_drops;
 pub mod pytket;
 pub mod replace_bools;
+pub mod replace_borrow_arrays;
 
 use derive_more::{Display, Error, From};
 use hugr::{
@@ -24,6 +25,7 @@ use hugr::{
 };
 use lower_drops::LowerDropsPass;
 use replace_bools::{ReplaceBoolPass, ReplaceBoolPassError};
+use replace_borrow_arrays::{ReplaceBorrowArrayPass, ReplaceBorrowArrayPassError};
 use tket::TketOp;
 
 use extension::{
@@ -44,6 +46,7 @@ pub struct QSystemPass {
     monomorphize: bool,
     force_order: bool,
     lazify: bool,
+    lower_borrow_arrays: bool,
 }
 
 impl Default for QSystemPass {
@@ -53,6 +56,7 @@ impl Default for QSystemPass {
             monomorphize: true,
             force_order: true,
             lazify: true,
+            lower_borrow_arrays: true,
         }
     }
 }
@@ -63,6 +67,8 @@ impl Default for QSystemPass {
 pub enum QSystemPassError<N = Node> {
     /// An error from the component [ReplaceBoolPass].
     ReplaceBoolError(ReplaceBoolPassError<N>),
+    /// An error from the component [ReplaceBorrowArrayPass].
+    ReplaceBorrowArrayError(ReplaceBorrowArrayPassError<N>),
     /// An error from the component [force_order()] pass.
     ForceOrderError(HugrError),
     /// An error from the component [LowerTketToQSystemPass] pass.
@@ -116,12 +122,19 @@ impl QSystemPass {
         }
 
         self.lower_tk2().run(hugr)?;
+        // Only has an effect if there are borrow arrays - should be run before
+        // lazification and drop lowering as no copy discard handler currently exists
+        // for BorrowArrays.
+        if self.lower_borrow_arrays {
+            self.replace_borrow_arrays().run(hugr)?;
+        }
         if self.lazify {
             self.replace_bools().run(hugr)?;
         }
-
         // We expect any Hugr will have *either* drop ops, or ValueArrays (without drops),
         // so only one of these passes will do anything; the order is thus immaterial.
+        // Drop should come after borrow array replacement so that we don't require a
+        // copy/discard handler for borrow arrays + avoid lowering the discard function.
         self.lower_drops().run(hugr)?;
         self.linearize_arrays().run(hugr)?;
 
@@ -190,6 +203,10 @@ impl QSystemPass {
         ReplaceBoolPass
     }
 
+    fn replace_borrow_arrays(&self) -> ReplaceBorrowArrayPass {
+        ReplaceBorrowArrayPass
+    }
+
     fn constant_fold(&self) -> ConstantFoldPass {
         ConstantFoldPass::default()
     }

tket-qsystem/src/lower_drops.rs

@@ -2,10 +2,14 @@
 use hugr::algorithms::replace_types::{NodeTemplate, ReplaceTypesError, ReplacementOptions};
 use hugr::algorithms::{ComposablePass, ReplaceTypes};
 use hugr::builder::{Container, DFGBuilder};
+use hugr::extension::prelude::bool_t;
+use hugr::extension::simple_op::MakeRegisteredOp;
 use hugr::types::{Signature, Term};
 use hugr::{hugr::hugrmut::HugrMut, Node};
 use tket::extension::guppy::{DROP_OP_NAME, GUPPY_EXTENSION};
 
+use crate::extension::futures::{future_type, FutureOp, FutureOpDef};
+
 /// A pass that lowers "drop" ops from [GUPPY_EXTENSION]
 #[derive(Default, Debug, Clone)]
 pub struct LowerDropsPass;
@@ -18,6 +22,30 @@ impl<H: HugrMut<Node = Node>> ComposablePass<H> for LowerDropsPass {
 
     fn run(&self, hugr: &mut H) -> Result<Self::Result, Self::Error> {
         let mut rt = ReplaceTypes::default();
+
+        // future(bool) is not in the default linearizer handler so we add it here.
+        // TODO: Create ReplaceTypes with future(bool) linearized by default to avoid
+        // code duplication with ReplaceBools pass.
+        let dup_op = FutureOp {
+            op: FutureOpDef::Dup,
+            typ: bool_t(),
+        }
+        .to_extension_op()
+        .unwrap();
+        let free_op = FutureOp {
+            op: FutureOpDef::Free,
+            typ: bool_t(),
+        }
+        .to_extension_op()
+        .unwrap();
+        rt.linearizer()
+            .register_simple(
+                future_type(bool_t()).as_extension().unwrap().clone(),
+                NodeTemplate::SingleOp(dup_op.into()),
+                NodeTemplate::SingleOp(free_op.into()),
+            )
+            .unwrap();
+
         rt.replace_parametrized_op_with(
             GUPPY_EXTENSION.get_op(DROP_OP_NAME.as_str()).unwrap(),
             |targs| {

tket-qsystem/src/replace_bools.rs

@@ -5,27 +5,34 @@ mod static_array;
 use derive_more::{Display, Error, From};
 use hugr::{
     algorithms::{
+        ensure_no_nonlocal_edges,
         non_local::FindNonLocalEdgesError,
-        replace_types::{NodeTemplate, ReplaceTypesError},
+        replace_types::{NodeTemplate, ReplaceTypesError, ReplacementOptions},
         ComposablePass, ReplaceTypes,
     },
     builder::{
-        inout_sig, BuildHandle, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer,
-        SubContainer,
+        inout_sig, BuildHandle, Container, DFGBuilder, Dataflow, DataflowHugr,
+        DataflowSubContainer, SubContainer,
     },
     extension::{
         prelude::{bool_t, qb_t},
         simple_op::MakeRegisteredOp,
     },
     hugr::hugrmut::HugrMut,
-    ops::{handle::ConditionalID, Tag, Value},
-    std_extensions::logic::LogicOp,
+    ops::{handle::ConditionalID, ExtensionOp, Tag, Value},
+    std_extensions::{
+        collections::array::{self, array_type, ARRAY_CLONE_OP_ID, ARRAY_DISCARD_OP_ID},
+        logic::LogicOp,
+    },
     types::{SumType, Type},
-    Hugr, Node, Wire,
+    Hugr, HugrView, Node, Wire,
 };
 use static_array::{ReplaceStaticArrayBoolPass, ReplaceStaticArrayBoolPassError};
 use tket::{
-    extension::bool::{bool_type, BoolOp, ConstBool},
+    extension::{
+        bool::{bool_type, BoolOp, ConstBool},
+        guppy::{DROP_OP_NAME, GUPPY_EXTENSION},
+    },
     TketOp,
 };
 
@@ -71,8 +78,7 @@ impl<H: HugrMut<Node = Node>> ComposablePass<H> for ReplaceBoolPass {
     type Result = ();
 
     fn run(&self, hugr: &mut H) -> Result<(), Self::Error> {
-        // TODO uncomment once https://github.com/CQCL/hugr/issues/1234 is complete
-        // ensure_no_nonlocal_edges(hugr)?;
+        ensure_no_nonlocal_edges(hugr)?;
         ReplaceStaticArrayBoolPass::default().run(hugr)?;
         let lowerer = lowerer();
         lowerer.run(hugr)?;
@@ -213,6 +219,20 @@ fn measure_reset_dest() -> NodeTemplate {
     NodeTemplate::CompoundOp(Box::new(h))
 }
 
+fn array_clone_dest(size: u64, elem_ty: Type) -> NodeTemplate {
+    let array_ty = array_type(size, elem_ty.clone());
+    let mut dfb = DFGBuilder::new(inout_sig(
+        vec![array_ty.clone()],
+        vec![array_ty.clone(), array_ty],
+    ))
+    .unwrap();
+    let mut h = std::mem::take(dfb.hugr_mut());
+    let [inp, outp] = h.get_io(h.entrypoint()).unwrap();
+    h.connect(inp, 0, outp, 0);
+    h.connect(inp, 0, outp, 1);
+    NodeTemplate::CompoundOp(Box::new(h))
+}
+
 /// The configuration used for replacing tket.bool extension types and ops.
 fn lowerer() -> ReplaceTypes {
     let mut lw = ReplaceTypes::default();
@@ -277,6 +297,44 @@ fn lowerer() -> ReplaceTypes {
     lw.replace_op(&qsystem_measure, measure_dest());
     lw.replace_op(&qsystem_measure_reset, measure_reset_dest());
 
+    // Replace array ops with copyable bounds with DFGs that the linearizer can act on
+    // now that the elements are no longer copyable.
+    lw.replace_parametrized_op_with(
+        array::EXTENSION.get_op(ARRAY_CLONE_OP_ID.as_str()).unwrap(),
+        move |args| {
+            let [size, elem_ty] = args else {
+                unreachable!()
+            };
+            let size = size.as_nat().unwrap();
+            let elem_ty = elem_ty.as_runtime().unwrap();
+            (!elem_ty.copyable()).then(|| array_clone_dest(size, elem_ty))
+        },
+        ReplacementOptions::default().with_linearization(true),
+    );
+
+    lw.replace_parametrized_op(
+        array::EXTENSION
+            .get_op(ARRAY_DISCARD_OP_ID.as_str())
+            .unwrap(),
+        move |args| {
+            let [size, elem_ty] = args else {
+                unreachable!()
+            };
+            let size = size.as_nat().unwrap();
+            let elem_ty = elem_ty.as_runtime().unwrap();
+            if elem_ty.copyable() {
+                return None;
+            }
+            let drop_op_def = GUPPY_EXTENSION.get_op(DROP_OP_NAME.as_str()).unwrap();
+            let drop_op = ExtensionOp::new(
+                drop_op_def.clone(),
+                vec![array_type(size, elem_ty.clone()).into()],
+            )
+            .unwrap();
+            Some(NodeTemplate::SingleOp(drop_op.into()))
+        },
+    );
+
     lw
 }
 
@@ -286,6 +344,8 @@ mod test {
 
     use super::*;
     use hugr::ops::OpType;
+    use hugr::std_extensions::collections::array::ArrayOpBuilder;
+    use hugr::type_row;
     use hugr::{
         builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr},
         extension::prelude::qb_t,
@@ -437,4 +497,49 @@ mod test {
         let sig = h.signature(h.entrypoint()).unwrap();
         assert_eq!(sig.output(), &TypeRow::from(vec![qb_t(), bool_dest()]));
     }
+
+    #[test]
+    fn test_array_clone_bool() {
+        let elem_ty = bool_type();
+        let size = 4;
+        let arr_ty = array_type(size, elem_ty.clone());
+        let mut dfb = DFGBuilder::new(inout_sig(
+            vec![arr_ty.clone()],
+            vec![arr_ty.clone(), arr_ty.clone()],
+        ))
+        .unwrap();
+        let [arr_in] = dfb.input_wires_arr();
+        let (arr1, arr2) = dfb.add_array_clone(elem_ty, size, arr_in).unwrap();
+        let mut h = dfb.finish_hugr_with_outputs([arr1, arr2]).unwrap();
+
+        h.validate().unwrap();
+        let pass = ReplaceBoolPass;
+        pass.run(&mut h).unwrap();
+        h.validate().unwrap();
+
+        let sig = h.signature(h.entrypoint()).unwrap();
+        let bool_dest_ty = bool_dest();
+        let arr_dest_ty = array_type(size, bool_dest_ty);
+        assert_eq!(sig.input(), &TypeRow::from(vec![arr_dest_ty.clone()]));
+        assert_eq!(
+            sig.output(),
+            &TypeRow::from(vec![arr_dest_ty.clone(), arr_dest_ty])
+        );
+    }
+
+    #[test]
+    fn test_array_discard_bool() {
+        let elem_ty = bool_type();
+        let size = 4;
+        let arr_ty = array_type(size, elem_ty.clone());
+        let mut dfb = DFGBuilder::new(inout_sig(vec![arr_ty.clone()], type_row![])).unwrap();
+        let [arr_in] = dfb.input_wires_arr();
+        dfb.add_array_discard(elem_ty, size, arr_in).unwrap();
+        let mut h = dfb.finish_hugr_with_outputs([]).unwrap();
+
+        h.validate().unwrap();
+        let pass = ReplaceBoolPass;
+        pass.run(&mut h).unwrap();
+        h.validate().unwrap();
+    }
 }