feat: LLVM lowering for borrow arrays using bitmasks (#2574)

Closes #2551

Substantial duplication between existing
hugr-llvm/src/extension/collections/array.rs and new ...borrow_array.rs.

---------

Co-authored-by: Seyon Sivarajah <[email protected]>
Co-authored-by: Alan Lawrence <[email protected]>

Author: 3 people

hugr-core/src/std_extensions/collections/borrow_array.rs

@@ -622,7 +622,7 @@ pub trait BArrayOpBuilder: GenericArrayOpBuilder {
         elem_ty: Type,
         input: Wire,
     ) -> Result<(), BuildError> {
-        self.add_generic_array_discard_empty::<Array>(elem_ty, input)
+        self.add_generic_array_discard_empty::<BorrowArray>(elem_ty, input)
     }
 
     /// Adds a borrow array borrow operation to the dataflow graph.

hugr-llvm/src/emit/func.rs

@@ -11,11 +11,11 @@ use inkwell::{
     basic_block::BasicBlock,
     builder::Builder,
     context::Context,
-    module::Module,
+    module::{Linkage, Module},
     types::{BasicType, BasicTypeEnum, FunctionType},
-    values::{BasicValueEnum, FunctionValue, GlobalValue, IntValue},
+    values::{BasicValue, BasicValueEnum, FunctionValue, GlobalValue, IntValue},
 };
-use itertools::zip_eq;
+use itertools::{Itertools, zip_eq};
 
 use crate::types::{HugrFuncType, HugrSumType, HugrType, TypingSession};
 use crate::{custom::CodegenExtsMap, types::LLVMSumType, utils::fat::FatNode};
@@ -357,6 +357,70 @@ pub fn build_ok_or_else<'c, H: HugrView<Node = Node>>(
     let either = builder.build_select(is_ok, right, left, "")?;
     Ok(either)
 }
+/// Helper to outline LLVM IR into a function call instead of inlining it every time.
+///
+/// The first time this helper is called with a given function name, a function is built
+/// using the provided closure. Future invocations with the same name will just emit calls
+/// to this function.
+///
+/// The return type is specified by `ret_type`, and if `Some` then the closure must return
+/// a value of that type, which will be returned from the function. Otherwise, the function
+/// will return void.
+pub fn get_or_make_function<'c, H: HugrView<Node = Node>, const N: usize>(
+    ctx: &mut EmitFuncContext<'c, '_, H>,
+    func_name: &str,
+    args: [BasicValueEnum<'c>; N],
+    ret_type: Option<BasicTypeEnum<'c>>,
+    go: impl FnOnce(
+        &mut EmitFuncContext<'c, '_, H>,
+        [BasicValueEnum<'c>; N],
+    ) -> Result<Option<BasicValueEnum<'c>>>,
+) -> Result<Option<BasicValueEnum<'c>>> {
+    let func = match ctx.get_current_module().get_function(func_name) {
+        Some(func) => func,
+        None => {
+            let arg_tys = args.iter().map(|v| v.get_type().into()).collect_vec();
+            let sig = match ret_type {
+                Some(ret_ty) => ret_ty.fn_type(&arg_tys, false),
+                None => ctx.iw_context().void_type().fn_type(&arg_tys, false),
+            };
+            let func =
+                ctx.get_current_module()
+                    .add_function(func_name, sig, Some(Linkage::Internal));
+            let bb = ctx.iw_context().append_basic_block(func, "");
+            let args = (0..N)
+                .map(|i| func.get_nth_param(i as u32).unwrap())
+                .collect_array()
+                .unwrap();
+
+            let curr_bb = ctx.builder().get_insert_block().unwrap();
+            let curr_func = ctx.func;
+
+            ctx.builder().position_at_end(bb);
+            ctx.func = func;
+            let ret_val = go(ctx, args)?;
+            if ctx
+                .builder()
+                .get_insert_block()
+                .unwrap()
+                .get_terminator()
+                .is_none()
+            {
+                ctx.builder()
+                    .build_return(ret_val.as_ref().map::<&dyn BasicValue, _>(|v| v))?;
+            }
+
+            ctx.builder().position_at_end(curr_bb);
+            ctx.func = curr_func;
+            func
+        }
+    };
+    let call_site =
+        ctx.builder()
+            .build_call(func, &args.iter().map(|&a| a.into()).collect_vec(), "")?;
+    let result = call_site.try_as_basic_value().left();
+    Ok(result)
+}
 
 #[cfg(test)]
 mod tests {

hugr-llvm/src/extension/collections.rs

@@ -1,6 +1,7 @@
 //! Emission logic for collections.
 
 pub mod array;
+pub mod borrow_array;
 pub mod list;
 pub mod stack_array;
 pub mod static_array;