diff --git a/hugr-llvm/src/extension/collections/borrow_array.rs b/hugr-llvm/src/extension/collections/borrow_array.rs index b59d9bac8..87021068e 100644 --- a/hugr-llvm/src/extension/collections/borrow_array.rs +++ b/hugr-llvm/src/extension/collections/borrow_array.rs @@ -1660,6 +1660,7 @@ pub fn emit_to_array_op<'c, H: HugrView>( offset, } = decompose_barray_fat_pointer(ctx.builder(), barray_v)?; build_none_borrowed_check(ccg, ctx, mask_ptr, offset, op.size)?; + ccg.emit_free_array(ctx, mask_ptr)?; Ok(build_array_fat_pointer(ctx, elems_ptr, offset)?.into()) } @@ -3120,6 +3121,106 @@ mod test { assert_eq!(&exec_ctx.exec_hugr_panicking(hugr, "main"), msg); } + static TO_ARRAY_ALLOCS: Mutex = Mutex::new(AllocationAccounting { + outstanding: Vec::new(), + allocated: 0, + invalid_frees: 0, + }); + + unsafe extern "C" fn tracked_alloc(size: usize) -> *mut c_void { + let ptr = unsafe { malloc(size) }; + let mut accounting = TO_ARRAY_ALLOCS.lock().unwrap(); + accounting.outstanding.push(ptr as usize); + accounting.allocated += 1; + ptr + } + + unsafe extern "C" fn tracked_free(ptr: *mut c_void) { + let mut accounting = TO_ARRAY_ALLOCS.lock().unwrap(); + if let Some(i) = accounting + .outstanding + .iter() + .position(|&p| p == ptr as usize) + { + accounting.outstanding.swap_remove(i); + unsafe { free(ptr) }; + } else { + accounting.invalid_frees += 1; + } + } + + #[rstest] + fn exec_to_array_frees_mask(mut exec_ctx: TestContext) { + use crate::utils::fat::FatExt as _; + + let int_ty = int_type(6); + let registry = ExtensionRegistry::new([ + int_types::EXTENSION.to_owned(), + prelude::PRELUDE.to_owned(), + borrow_array::EXTENSION.to_owned(), + hugr_core::std_extensions::collections::array::EXTENSION.to_owned(), + ]); + let hugr = SimpleHugrConfig::new() + .with_extensions(registry) + .with_outs([int_ty.clone()]) + .finish(|mut b| { + let value = b.add_load_value(ConstInt::new_u(6, 7).unwrap()); + let barray = b.add_new_borrow_array(int_ty.clone(), [value]).unwrap(); + let array = b + .add_dataflow_op(BArrayToArray::new(int_ty.clone(), 1), [barray]) + .unwrap() + .out_wire(0); + let index = b.add_load_value(ConstUsize::new(0)); + let (value, array) = b.add_array_get(int_ty.clone(), 1, array, index).unwrap(); + let [value] = b + .build_unwrap_sum(1, option_type([int_ty.clone()]), value) + .unwrap(); + b.add_array_discard(int_ty.clone(), 1, array).unwrap(); + b.finish_hugr_with_outputs([value]).unwrap() + }); + exec_ctx.add_extensions(|cge| { + cge.add_default_prelude_extensions() + .add_default_borrow_array_extensions(DefaultPreludeCodegen) + .add_default_array_extensions() + .add_default_int_extensions() + }); + let emission = Emission::emit_hugr( + hugr.fat_root().unwrap(), + exec_ctx.get_emit_hugr(), + TEST_EMIT_DEBUG, + ) + .unwrap(); + emission.verify().unwrap(); + let module = emission.module(); + module + .get_function("main") + .unwrap() + .set_linkage(inkwell::module::Linkage::External); + let engine = module + .create_jit_execution_engine(inkwell::OptimizationLevel::None) + .unwrap(); + engine.add_global_mapping( + &module.get_function("malloc").unwrap(), + tracked_alloc as *const () as usize, + ); + engine.add_global_mapping( + &module.get_function("free").unwrap(), + tracked_free as *const () as usize, + ); + let run: unsafe extern "C" fn() -> u64 = + unsafe { engine.get_function("main").unwrap().into_raw() }; + for _ in 0..100 { + assert_eq!(unsafe { run() }, 7); + } + let accounting = TO_ARRAY_ALLOCS.lock().unwrap(); + assert_eq!(accounting.allocated, 200); + assert_eq!(accounting.invalid_frees, 0); + assert!( + accounting.outstanding.is_empty(), + "borrow-array masks leaked" + ); + } + #[rstest] fn exec_is_borrowed_basic(mut exec_ctx: TestContext) { // We build a HUGR that: