refactor: allow non-void get_or_make_function (#2609)

Author: ss2165

hugr-llvm/src/emit/func.rs

@@ -357,23 +357,32 @@ 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 the closure must return a value of that type.
+/// If `ret_type` is `None`, the function is assumed to 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],
-    go: impl FnOnce(&mut EmitFuncContext<'c, '_, H>, [BasicValueEnum<'c>; N]) -> Result<()>,
-) -> Result<()> {
+    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 = ctx.iw_context().void_type().fn_type(&arg_tys, false);
+            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));
@@ -388,25 +397,30 @@ pub fn get_or_make_function<'c, H: HugrView<Node = Node>, const N: usize>(
 
             ctx.builder().position_at_end(bb);
             ctx.func = func;
-            go(ctx, args)?;
+            let ret_val = go(ctx, args)?;
             if ctx
                 .builder()
                 .get_insert_block()
                 .unwrap()
                 .get_terminator()
                 .is_none()
             {
-                ctx.builder().build_return(None)?;
+                match ret_val {
+                    Some(ref v) => ctx.builder().build_return(Some(v))?,
+                    None => ctx.builder().build_return(None)?,
+                };
             }
 
             ctx.builder().position_at_end(curr_bb);
             ctx.func = curr_func;
             func
         }
     };
-    ctx.builder()
-        .build_call(func, &args.iter().map(|&a| a.into()).collect_vec(), "")?;
-    Ok(())
+    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)]

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

@@ -679,6 +679,7 @@ fn build_mask_flip<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [mask_ptr.into(), idx.into()],
+        None,
         |ctx, [mask_ptr, idx]| {
             let mask_ptr = mask_ptr.into_pointer_value();
             let idx = idx.into_int_value();
@@ -694,9 +695,10 @@ fn build_mask_flip<'c, H: HugrView<Node = Node>>(
             let update = builder.build_left_shift(usize_t.const_int(1, false), block_offset, "")?;
             let block = builder.build_xor(block, update, "")?;
             builder.build_store(block_addr, block)?;
-            Ok(())
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Emits a check that a specific array element has not already been borrowed.
@@ -712,6 +714,7 @@ pub fn build_idx_not_borrowed_check<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [mask_ptr.into(), idx.into()],
+        None,
         |ctx, [mask_ptr, idx]| {
             // Emit panic if borrow-bit is set
             inspect_mask_idx(
@@ -726,9 +729,11 @@ pub fn build_idx_not_borrowed_check<'c, H: HugrView<Node = Node>>(
                     ctx.builder().build_unreachable()?;
                     Ok(())
                 },
-            )
+            )?;
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Emits a check that a specific array index is free.
@@ -744,6 +749,7 @@ pub fn build_idx_free_check<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [mask_ptr.into(), idx.into()],
+        None,
         |ctx, [mask_ptr, idx]| {
             // Emit panic if borrow-bit is not set
             inspect_mask_idx(
@@ -758,9 +764,11 @@ pub fn build_idx_free_check<'c, H: HugrView<Node = Node>>(
                     Ok(())
                 },
                 |_| Ok(()),
-            )
+            )?;
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Emits a check that no array elements have been borrowed.
@@ -782,6 +790,7 @@ pub fn build_none_borrowed_check<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [mask_ptr.into(), offset.into(), size.into()],
+        None,
         |ctx, [mask_ptr, offset, size]| {
             let mask_ptr = mask_ptr.into_pointer_value();
             let offset = offset.into_int_value();
@@ -806,9 +815,11 @@ pub fn build_none_borrowed_check<'c, H: HugrView<Node = Node>>(
                 builder.build_conditional_branch(cond, none_borrowed_bb, some_borrowed_bb)?;
                 builder.position_at_end(none_borrowed_bb);
                 Ok(())
-            })
+            })?;
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Emits a check that all array elements have been borrowed.
@@ -830,6 +841,7 @@ pub fn build_all_borrowed_check<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [mask_ptr.into(), offset.into(), size.into()],
+        None,
         |ctx, [mask_ptr, offset, size]| {
             let mask_ptr = mask_ptr.into_pointer_value();
             let offset = offset.into_int_value();
@@ -858,9 +870,11 @@ pub fn build_all_borrowed_check<'c, H: HugrView<Node = Node>>(
                 builder.build_conditional_branch(cond, all_borrowed_bb, some_not_borrowed)?;
                 builder.position_at_end(all_borrowed_bb);
                 Ok(())
-            })
+            })?;
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Emits a check that a specific array element has not already been borrowed.
@@ -877,6 +891,7 @@ pub fn build_bounds_check<'c, H: HugrView<Node = Node>>(
         ctx,
         FUNC_NAME,
         [size.into(), idx.into()],
+        None,
         |ctx, [size, idx]| {
             let size = size.into_int_value();
             let idx = idx.into_int_value();
@@ -894,9 +909,10 @@ pub fn build_bounds_check<'c, H: HugrView<Node = Node>>(
             ctx.builder()
                 .build_conditional_branch(in_bounds, ok_bb, err_bb)?;
             ctx.builder().position_at_end(ok_bb);
-            Ok(())
+            Ok(None)
         },
-    )
+    )?;
+    Ok(())
 }
 
 /// Helper function to build a loop that repeats for a given number of iterations.