fix!: Change array result ops signature to return array result (#888)

Needed for CQCL/guppylang#981

BREAKING CHANGE: `ResultOpDef::ArrBool`, `ResultOpDef::ArrInt`,
`ResultOpDef::ArrUInt` and `ResultOpDef::ArrF64` signatures now return
array results

Author: tatiana-s

tket2-exts/src/tket2_exts/data/tket2/result.json

@@ -44,7 +44,32 @@
               "bound": "A"
             }
           ],
-          "output": []
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "collections.array",
+              "id": "array",
+              "args": [
+                {
+                  "tya": "Variable",
+                  "idx": 1,
+                  "cached_decl": {
+                    "tp": "BoundedNat",
+                    "bound": null
+                  }
+                },
+                {
+                  "tya": "Type",
+                  "ty": {
+                    "t": "Sum",
+                    "s": "Unit",
+                    "size": 2
+                  }
+                }
+              ],
+              "bound": "A"
+            }
+          ]
         }
       },
       "binary": false
@@ -92,7 +117,34 @@
               "bound": "A"
             }
           ],
-          "output": []
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "collections.array",
+              "id": "array",
+              "args": [
+                {
+                  "tya": "Variable",
+                  "idx": 1,
+                  "cached_decl": {
+                    "tp": "BoundedNat",
+                    "bound": null
+                  }
+                },
+                {
+                  "tya": "Type",
+                  "ty": {
+                    "t": "Opaque",
+                    "extension": "arithmetic.float.types",
+                    "id": "float64",
+                    "args": [],
+                    "bound": "C"
+                  }
+                }
+              ],
+              "bound": "A"
+            }
+          ]
         }
       },
       "binary": false
@@ -153,7 +205,43 @@
               "bound": "A"
             }
           ],
-          "output": []
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "collections.array",
+              "id": "array",
+              "args": [
+                {
+                  "tya": "Variable",
+                  "idx": 1,
+                  "cached_decl": {
+                    "tp": "BoundedNat",
+                    "bound": null
+                  }
+                },
+                {
+                  "tya": "Type",
+                  "ty": {
+                    "t": "Opaque",
+                    "extension": "arithmetic.int.types",
+                    "id": "int",
+                    "args": [
+                      {
+                        "tya": "Variable",
+                        "idx": 2,
+                        "cached_decl": {
+                          "tp": "BoundedNat",
+                          "bound": 7
+                        }
+                      }
+                    ],
+                    "bound": "C"
+                  }
+                }
+              ],
+              "bound": "A"
+            }
+          ]
         }
       },
       "binary": false
@@ -214,7 +302,43 @@
               "bound": "A"
             }
           ],
-          "output": []
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "collections.array",
+              "id": "array",
+              "args": [
+                {
+                  "tya": "Variable",
+                  "idx": 1,
+                  "cached_decl": {
+                    "tp": "BoundedNat",
+                    "bound": null
+                  }
+                },
+                {
+                  "tya": "Type",
+                  "ty": {
+                    "t": "Opaque",
+                    "extension": "arithmetic.int.types",
+                    "id": "int",
+                    "args": [
+                      {
+                        "tya": "Variable",
+                        "idx": 2,
+                        "cached_decl": {
+                          "tp": "BoundedNat",
+                          "bound": 7
+                        }
+                      }
+                    ],
+                    "bound": "C"
+                  }
+                }
+              ],
+              "bound": "A"
+            }
+          ]
         }
       },
       "binary": false

tket2-hseries/src/extension/result.rs

@@ -161,11 +161,20 @@ impl ResultOpDef {
     }
 
     fn result_signature(&self) -> SignatureFunc {
-        PolyFuncType::new(
-            [vec![TypeParam::String], self.type_params()].concat(),
-            Signature::new(self.arg_type(), type_row![]),
+        let sig = if self.is_array_result_op() {
+            // Do not consume input arrays to allow them to be not copyable.
+            Signature::new(self.arg_type(), self.arg_type())
+        } else {
+            Signature::new(self.arg_type(), type_row![])
+        };
+        PolyFuncType::new([vec![TypeParam::String], self.type_params()].concat(), sig).into()
+    }
+
+    fn is_array_result_op(&self) -> bool {
+        matches!(
+            self,
+            Self::ArrBool | Self::ArrF64 | Self::ArrInt | Self::ArrUInt
         )
-        .into()
     }
 }
 
@@ -390,19 +399,32 @@ impl TryFrom<&OpType> for ResultOpDef {
 pub trait ResultOpBuilder: Dataflow {
     /// Add a "tket2.result" op.
     fn add_result(&mut self, result_wire: Wire, op: ResultOp) -> Result<(), BuildError> {
+        debug_assert!(!op.result_op.is_array_result_op());
         let handle = self.add_dataflow_op(op, [result_wire])?;
 
         debug_assert_eq!(handle.outputs().len(), 0);
         Ok(())
     }
+
+    /// Add a "tket2.result" op with array result type.
+    fn add_array_result(
+        &mut self,
+        result_wire: Wire,
+        op: ResultOp,
+    ) -> Result<[Wire; 1], BuildError> {
+        debug_assert!(op.result_op.is_array_result_op());
+        Ok(self
+            .add_dataflow_op(op.clone(), [result_wire])?
+            .outputs_arr())
+    }
 }
 
 impl<D: Dataflow> ResultOpBuilder for D {}
 
 #[cfg(test)]
 pub(crate) mod test {
     use cool_asserts::assert_matches;
-    use hugr::types::Signature;
+    use hugr::types::{NoRV, Signature, TypeBase};
     use hugr::HugrView;
     use hugr::{
         builder::{Dataflow, DataflowHugr, FunctionBuilder},
@@ -436,17 +458,15 @@ pub(crate) mod test {
             INT_TYPES[5].clone(),
             INT_TYPES[6].clone(),
         ];
-        let in_row = [
-            in_row.clone(),
-            in_row
-                .into_iter()
-                .map(|t| array_type(ARR_SIZE, t))
-                .collect(),
-        ]
-        .concat();
+        let arrs: Vec<TypeBase<NoRV>> = in_row
+            .clone()
+            .into_iter()
+            .map(|t| array_type(ARR_SIZE, t))
+            .collect();
+        let in_row = [in_row.clone(), arrs.clone()].concat();
         let hugr = {
             let mut func_builder =
-                FunctionBuilder::new("circuit", Signature::new(in_row, type_row![])).unwrap();
+                FunctionBuilder::new("circuit", Signature::new(in_row, arrs)).unwrap();
             let ops = [
                 ResultOp::new_bool("b"),
                 ResultOp::new_f64("f"),
@@ -474,13 +494,15 @@ pub(crate) mod test {
             for (w, op) in [b, f, i, u].iter().zip(ops.iter()) {
                 func_builder.add_result(*w, op.clone()).unwrap();
             }
+            let mut outputs = Vec::new();
             for (w, op) in [a_b, a_f, a_i, a_u].iter().zip(ops.iter()) {
-                func_builder
-                    .add_result(*w, op.clone().array_op(ARR_SIZE))
+                let [out_w] = func_builder
+                    .add_array_result(*w, op.clone().array_op(ARR_SIZE))
                     .unwrap();
+                outputs.push(out_w);
             }
 
-            func_builder.finish_hugr_with_outputs([]).unwrap()
+            func_builder.finish_hugr_with_outputs(outputs).unwrap()
         };
         assert_matches!(hugr.validate(), Ok(_));
     }

tket2-hseries/src/llvm/result.rs

@@ -154,7 +154,6 @@ impl<'c, H: HugrView<Node = Node>> ResultEmitter<'c, '_, '_, H> {
         tag_ptr: BasicValueEnum,
         tag_len: IntValue,
     ) -> Result<()> {
-        // TODO update to return array after https://github.com/CQCL/tket2/pull/888
         let ResultArgs::Array(_, length) = op.args else {
             bail!("Expected array argument")
         };
@@ -241,6 +240,8 @@ impl<'c, H: HugrView<Node = Node>> ResultEmitter<'c, '_, '_, H> {
                     .try_into()
                     .map_err(|_| anyhow!("result_arr_bool expects one input"))?;
                 self.build_print_array_call(val, op, &ElemType::Bool, tag_ptr, tag_len)?;
+                // Array results need to output the input value again.
+                return args.outputs.finish(self.builder(), [val]);
             }
             ResultOpDef::ArrInt => {
                 let (tag_ptr, tag_len) = self.generate_global_tag(&args, "INTARR:").unwrap();
@@ -249,6 +250,7 @@ impl<'c, H: HugrView<Node = Node>> ResultEmitter<'c, '_, '_, H> {
                     .try_into()
                     .map_err(|_| anyhow!("result_arr_int expects one input"))?;
                 self.build_print_array_call(val, op, &ElemType::Int, tag_ptr, tag_len)?;
+                return args.outputs.finish(self.builder(), [val]);
             }
             ResultOpDef::ArrUInt => {
                 let (tag_ptr, tag_len) = self.generate_global_tag(&args, "INTARR:").unwrap();
@@ -257,6 +259,7 @@ impl<'c, H: HugrView<Node = Node>> ResultEmitter<'c, '_, '_, H> {
                     .try_into()
                     .map_err(|_| anyhow!("result_arr_uint expects one input"))?;
                 self.build_print_array_call(val, op, &ElemType::Uint, tag_ptr, tag_len)?;
+                return args.outputs.finish(self.builder(), [val]);
             }
             ResultOpDef::ArrF64 => {
                 let (tag_ptr, tag_len) = self.generate_global_tag(&args, "FLOATARR:").unwrap();
@@ -265,6 +268,7 @@ impl<'c, H: HugrView<Node = Node>> ResultEmitter<'c, '_, '_, H> {
                     .try_into()
                     .map_err(|_| anyhow!("result_arr_float expects one input"))?;
                 self.build_print_array_call(val, op, &ElemType::Float, tag_ptr, tag_len)?;
+                return args.outputs.finish(self.builder(), [val]);
             }
         }
         args.outputs.finish(self.builder(), [])

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__result__test__emit_result_codegen@llvm14_5.snap

@@ -7,7 +7,7 @@ source_filename = "test_context"
 
 @res_test_arr_b.866EEC87.0 = private constant [27 x i8] c"\1AUSER:BOOLARR:test_arr_bool"
 
-define void @_hl.main.1([10 x i1] %0) {
+define [10 x i1] @_hl.main.1([10 x i1] %0) {
 alloca_block:
   br label %entry_block
 
@@ -31,7 +31,7 @@ entry_block:                                      ; preds = %alloca_block
   store i1* %3, i1** %mask_ptr, align 8
   %5 = load <{ i32, i32, i1*, i1* }>, <{ i32, i32, i1*, i1* }>* %out_arr_alloca, align 1
   call void @print_bool_arr(i8* getelementptr inbounds ([27 x i8], [27 x i8]* @res_test_arr_b.866EEC87.0, i32 0, i32 0), i64 %tag_len2, <{ i32, i32, i1*, i1* }>* %out_arr_alloca)
-  ret void
+  ret [10 x i1] %0
 }
 
 declare void @print_bool_arr(i8*, i64, <{ i32, i32, i1*, i1* }>*)

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__result__test__emit_result_codegen@llvm14_6.snap

@@ -7,7 +7,7 @@ source_filename = "test_context"
 
 @res_test_arr_i.DFD30452.0 = private constant [25 x i8] c"\18USER:INTARR:test_arr_int"
 
-define void @_hl.main.1([10 x i64] %0) {
+define [10 x i64] @_hl.main.1([10 x i64] %0) {
 alloca_block:
   br label %entry_block
 
@@ -31,7 +31,7 @@ entry_block:                                      ; preds = %alloca_block
   store i1* %3, i1** %mask_ptr, align 8
   %5 = load <{ i32, i32, i64*, i1* }>, <{ i32, i32, i64*, i1* }>* %out_arr_alloca, align 1
   call void @print_int_arr(i8* getelementptr inbounds ([25 x i8], [25 x i8]* @res_test_arr_i.DFD30452.0, i32 0, i32 0), i64 %tag_len2, <{ i32, i32, i64*, i1* }>* %out_arr_alloca)
-  ret void
+  ret [10 x i64] %0
 }
 
 declare void @print_int_arr(i8*, i64, <{ i32, i32, i64*, i1* }>*)

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__result__test__emit_result_codegen@llvm14_7.snap

@@ -7,7 +7,7 @@ source_filename = "test_context"
 
 @res_test_arr_u.3D1C515C.0 = private constant [26 x i8] c"\19USER:INTARR:test_arr_uint"
 
-define void @_hl.main.1([10 x i64] %0) {
+define [10 x i64] @_hl.main.1([10 x i64] %0) {
 alloca_block:
   br label %entry_block
 
@@ -31,7 +31,7 @@ entry_block:                                      ; preds = %alloca_block
   store i1* %3, i1** %mask_ptr, align 8
   %5 = load <{ i32, i32, i64*, i1* }>, <{ i32, i32, i64*, i1* }>* %out_arr_alloca, align 1
   call void @print_uint_arr(i8* getelementptr inbounds ([26 x i8], [26 x i8]* @res_test_arr_u.3D1C515C.0, i32 0, i32 0), i64 %tag_len2, <{ i32, i32, i64*, i1* }>* %out_arr_alloca)
-  ret void
+  ret [10 x i64] %0
 }
 
 declare void @print_uint_arr(i8*, i64, <{ i32, i32, i64*, i1* }>*)

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__result__test__emit_result_codegen@llvm14_8.snap

@@ -7,7 +7,7 @@ source_filename = "test_context"
 
 @res_test_arr_f.038B27BE.0 = private constant [27 x i8] c"\1AUSER:FLOATARR:test_arr_f64"
 
-define void @_hl.main.1([10 x double] %0) {
+define [10 x double] @_hl.main.1([10 x double] %0) {
 alloca_block:
   br label %entry_block
 
@@ -31,7 +31,7 @@ entry_block:                                      ; preds = %alloca_block
   store i1* %3, i1** %mask_ptr, align 8
   %5 = load <{ i32, i32, double*, i1* }>, <{ i32, i32, double*, i1* }>* %out_arr_alloca, align 1
   call void @print_float_arr(i8* getelementptr inbounds ([27 x i8], [27 x i8]* @res_test_arr_f.038B27BE.0, i32 0, i32 0), i64 %tag_len2, <{ i32, i32, double*, i1* }>* %out_arr_alloca)
-  ret void
+  ret [10 x double] %0
 }
 
 declare void @print_float_arr(i8*, i64, <{ i32, i32, double*, i1* }>*)