feat: add llvm codegen for is_borrowed

Author: ss2165

hugr-llvm/src/extension/collections/borrow_array.rs

@@ -572,6 +572,8 @@ enum MaskCheck {
     CheckNotBorrowed,
     /// Check the element is not borrowed, panicking if it is; then mark as borrowed.
     Borrow,
+    /// Check if the element is borrowed, returning a boolean. (Do not change the bit.)
+    IsBorrowed,
 }
 
 impl MaskCheck {
@@ -580,9 +582,17 @@ impl MaskCheck {
             MaskCheck::Return => "__barray_mask_return",
             MaskCheck::CheckNotBorrowed => "__barray_mask_check_not_borrowed",
             MaskCheck::Borrow => "__barray_mask_borrow",
+            MaskCheck::IsBorrowed => "__barray_mask_is_borrowed",
         }
     }
 
+    fn return_type<'c, H: HugrView<Node = Node>>(
+        &self,
+        ctx: &EmitFuncContext<'c, '_, H>,
+    ) -> Option<BasicTypeEnum<'c>> {
+        (*self == MaskCheck::IsBorrowed).then(|| ctx.iw_context().bool_type().into())
+    }
+
     /// Generate code to perform the check on the specified bit of the mask.
     /// (Does not check the index is in bounds.)
     fn emit<'c, H: HugrView<Node = Node>>(
@@ -591,12 +601,12 @@ impl MaskCheck {
         ctx: &mut EmitFuncContext<'c, '_, H>,
         mask_ptr: PointerValue<'c>,
         idx: IntValue<'c>,
-    ) -> Result<()> {
+    ) -> Result<Option<BasicValueEnum<'c>>> {
         get_or_make_function(
             ctx,
             self.func_name(),
             [mask_ptr.into(), idx.into()],
-            None,
+            self.return_type(ctx),
             |ctx, [mask_ptr, idx]| {
                 // Compute mask bitarray block index via `idx // BLOCK_SIZE`
                 let mask_ptr = mask_ptr.into_pointer_value();
@@ -613,10 +623,15 @@ impl MaskCheck {
                 let block_shifted = builder.build_right_shift(block, idx_in_block, false, "")?;
                 let bit =
                     builder.build_int_truncate(block_shifted, ctx.iw_context().bool_type(), "")?;
+                if *self == MaskCheck::IsBorrowed {
+                    // Just return the bit
+                    return Ok(Some(bit.as_basic_value_enum()));
+                }
                 let panic_bb = ctx.build_positioned_new_block("panic", None, |ctx, panic_bb| {
                     let err: &ConstError = match self {
                         MaskCheck::CheckNotBorrowed | MaskCheck::Borrow => &ERR_ALREADY_BORROWED,
                         MaskCheck::Return => &ERR_NOT_BORROWED,
+                        MaskCheck::IsBorrowed => unreachable!("handled above"),
                     };
                     let err_val = ctx.emit_custom_const(err).unwrap();
                     ccg.emit_panic(ctx, err_val)?;
@@ -641,13 +656,13 @@ impl MaskCheck {
                 let (if_borrowed, if_present) = match self {
                     MaskCheck::CheckNotBorrowed | MaskCheck::Borrow => (panic_bb, ok_bb),
                     MaskCheck::Return => (ok_bb, panic_bb),
+                    MaskCheck::IsBorrowed => unreachable!("handled above"),
                 };
                 ctx.builder()
                     .build_conditional_branch(bit, if_borrowed, if_present)?;
                 Ok(None)
             },
-        )?;
-        Ok(())
+        )
     }
 }
 
@@ -1570,6 +1585,22 @@ pub fn emit_barray_unsafe_op<'c, H: HugrView<Node = Node>>(
             let (_, array_v) = build_barray_alloc(ctx, ccg, elem_ty, size, true)?;
             outputs.finish(ctx.builder(), [array_v.into()])
         }
+        BArrayUnsafeOpDef::is_borrowed => {
+            let [array_v, index_v] = inputs
+                .try_into()
+                .map_err(|_| anyhow!("BArrayUnsafeOpDef::is_borrowed expects two arguments"))?;
+            let BArrayFatPtrComponents {
+                mask_ptr, offset, ..
+            } = decompose_barray_fat_pointer(builder, array_v)?;
+            let index_v = index_v.into_int_value();
+            build_bounds_check(ccg, ctx, size, index_v)?;
+            let offset_index_v = ctx.builder().build_int_add(index_v, offset, "")?;
+            // let bit = build_is_borrowed_check(ctx, mask_ptr, offset_index_v)?;
+            let bit = MaskCheck::IsBorrowed
+                .emit(ccg, ctx, mask_ptr, offset_index_v)?
+                .expect("IsBorrowed always returns a value");
+            outputs.finish(ctx.builder(), [bit.into(), array_v])
+        }
         _ => todo!(),
     }
 }
@@ -2634,7 +2665,7 @@ mod test {
         // - Pops specified numbers from the left to introduce an offset
         // - Converts it into a regular array
         // - Converts it back into a borrow array
-        // - Borrows alls elements, sums them up, and returns the sum
+        // - Borrows all elements, sums them up, and returns the sum
 
         let int_ty = int_type(6);
         let hugr = SimpleHugrConfig::new()
@@ -2908,4 +2939,77 @@ mod test {
         let msg = "Some array elements have been borrowed";
         assert_eq!(&exec_ctx.exec_hugr_panicking(hugr, "main"), msg);
     }
+
+    #[rstest]
+    fn exec_is_borrowed_basic(mut exec_ctx: TestContext) {
+        // We build a HUGR that:
+        // - Creates a borrow array [1,2,3]
+        // - Borrows index 1
+        // - Checks is_borrowed for indices 0, 1
+        // - Returns 1 if [false, true], else 0
+        let int_ty = int_type(6);
+        let size = 3;
+        let hugr = SimpleHugrConfig::new()
+            .with_outs(int_ty.clone())
+            .with_extensions(exec_registry())
+            .finish(|mut builder| {
+                let barray = borrow_array::BArrayValue::new(
+                    int_ty.clone(),
+                    (1..=3)
+                        .map(|i| ConstInt::new_u(6, i).unwrap().into())
+                        .collect_vec(),
+                );
+                let barray = builder.add_load_value(barray);
+                let idx1 = builder.add_load_value(ConstUsize::new(1));
+                let (_, barray) = builder
+                    .add_borrow_array_borrow(int_ty.clone(), size, barray, idx1)
+                    .unwrap();
+
+                let idx0 = builder.add_load_value(ConstUsize::new(0));
+                let (arr, b0_bools) =
+                    [idx0, idx1]
+                        .iter()
+                        .fold((barray, Vec::new()), |(arr, mut bools), idx| {
+                            let (b, arr) = builder
+                                .add_is_borrowed(int_ty.clone(), size, arr, *idx)
+                                .unwrap();
+                            bools.push(b);
+                            (arr, bools)
+                        });
+                let [b0, b1] = b0_bools.try_into().unwrap();
+
+                let b0 = builder.add_not(b0).unwrap(); // flip b0 to true
+                let and01 = builder.add_and(b0, b1).unwrap();
+                let one = builder.add_load_value(ConstInt::new_u(6, 1).unwrap());
+                let zero = builder.add_load_value(ConstInt::new_u(6, 0).unwrap());
+                let mut cond = builder
+                    .conditional_builder(
+                        ([type_row![], type_row![]], and01),
+                        [],
+                        int_ty.clone().into(),
+                    )
+                    .unwrap();
+                cond.case_builder(0)
+                    .unwrap()
+                    .finish_with_outputs([zero])
+                    .unwrap();
+                cond.case_builder(1)
+                    .unwrap()
+                    .finish_with_outputs([one])
+                    .unwrap();
+                let out = cond.finish_sub_container().unwrap().out_wire(0);
+                builder
+                    .add_borrow_array_discard(int_ty.clone(), size, arr)
+                    .unwrap();
+                builder.finish_hugr_with_outputs([out]).unwrap()
+            });
+
+        exec_ctx.add_extensions(|cge| {
+            cge.add_default_prelude_extensions()
+                .add_logic_extensions()
+                .add_default_borrow_array_extensions(DefaultPreludeCodegen)
+                .add_default_int_extensions()
+        });
+        assert_eq!(1, exec_ctx.exec_hugr_u64(hugr, "main"));
+    }
 }