Break build_mask_padding up into two calls to build_mask_padding1

Author: acl-cqc

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

@@ -586,11 +586,25 @@ fn check_all_mask_eq<'c, H: HugrView<Node = Node>>(
     expected: bool,
     err: &ConstError,
 ) -> Result<()> {
-    build_mask_padding(ctx, mask_info, expected)?;
+    let end_idx = ctx
+        .builder()
+        .build_int_add(mask_info.offset, mask_info.size, "")?;
+    build_mask_padding1(
+        ctx,
+        mask_info.mask_ptr,
+        mask_info.offset,
+        expected,
+        PaddingDirection::First,
+    )?;
+    build_mask_padding1(
+        ctx,
+        mask_info.mask_ptr,
+        end_idx,
+        expected,
+        PaddingDirection::Last,
+    )?;
 
     let builder = ctx.builder();
-    let end_idx = builder.build_int_add(mask_info.offset, mask_info.size, "")?;
-
     let usize_t = usize_ty(&ctx.typing_session());
     let expected_val = if expected {
         usize_t.const_all_ones()
@@ -624,58 +638,49 @@ fn check_all_mask_eq<'c, H: HugrView<Node = Node>>(
     })
 }
 
-/// Emits instructions to update the mask, overwriting unused bits with a value.
-fn build_mask_padding<'c, H: HugrView<Node = Node>>(
+#[derive(Copy, Clone, Debug)]
+enum PaddingDirection {
+    First,
+    Last,
+}
+
+/// Emits instructions to destructively update the first or last block of the mask,
+/// given `idx` the first or last element within the mask, overwriting unused bits with `value`.
+fn build_mask_padding1<'c, H: HugrView<Node = Node>>(
     ctx: &mut EmitFuncContext<'c, '_, H>,
-    info: &MaskInfo<'c>,
+    mask_ptr: PointerValue<'c>,
+    idx: IntValue<'c>,
     value: bool,
+    direction: PaddingDirection,
 ) -> Result<()> {
-    let MaskInfo {
-        mask_ptr,
-        offset,
-        size,
-    } = *info;
     let builder = ctx.builder();
     let usize_t = usize_ty(&ctx.typing_session());
     let block_size = usize_t.const_int(usize_t.get_bit_width() as u64, false);
 
-    // Find the first and last blocks that contain some used bits
-    let lst_idx = builder.build_int_add(offset, size, "")?;
-    let fst_block_idx = builder.build_int_unsigned_div(offset, block_size, "")?;
-    let lst_block_idx = builder.build_int_unsigned_div(lst_idx, block_size, "")?;
-    let fst_block_addr = unsafe { builder.build_in_bounds_gep(mask_ptr, &[fst_block_idx], "")? };
-    let lst_block_addr = unsafe { builder.build_in_bounds_gep(mask_ptr, &[lst_block_idx], "")? };
-    let fst_block = builder.build_load(fst_block_addr, "")?.into_int_value();
-    let lst_block = builder.build_load(lst_block_addr, "")?.into_int_value();
+    let block_idx = builder.build_int_unsigned_div(idx, block_size, "")?;
+    let block_addr = unsafe { builder.build_in_bounds_gep(mask_ptr, &[block_idx], "")? };
+    let block = builder.build_load(block_addr, "")?.into_int_value();
 
-    // Pad out the unused bits in the first block
-    let ones = usize_t.const_all_ones();
-    let fst_block_unused = builder.build_int_unsigned_rem(offset, block_size, "")?;
-    let fst_block_used = builder.build_int_sub(block_size, fst_block_unused, "")?;
-    let new_fst_block = if value {
-        // Pad with ones
-        let pad = builder.build_right_shift(ones, fst_block_used, false, "")?;
-        builder.build_or(fst_block, pad, "")?
-    } else {
-        // Pad with zeros
-        let pad = builder.build_left_shift(ones, fst_block_unused, "")?;
-        builder.build_and(fst_block, pad, "")?
-    };
-    builder.build_store(fst_block_addr, new_fst_block)?;
+    let idx_in_block = builder.build_int_unsigned_rem(idx, block_size, "")?;
+    let idx_from_end = builder.build_int_sub(block_size, idx_in_block, "")?;
 
-    // Pad out the unused bits in the last block
-    let lst_block_used = builder.build_int_unsigned_rem(lst_idx, block_size, "")?;
-    let lst_block_unused = builder.build_int_sub(block_size, lst_block_used, "")?;
-    let new_lst_block = if value {
+    let ones = usize_t.const_all_ones();
+    let new_block = if value {
         // Pad with ones
-        let pad = builder.build_left_shift(ones, lst_block_used, "")?;
-        builder.build_or(lst_block, pad, "")?
+        let pad = match direction {
+            PaddingDirection::First => builder.build_right_shift(ones, idx_from_end, false, "")?,
+            PaddingDirection::Last => builder.build_left_shift(ones, idx_in_block, "")?,
+        };
+        builder.build_or(block, pad, "")?
     } else {
         // Pad with zeros
-        let pad = builder.build_right_shift(ones, lst_block_unused, false, "")?;
-        builder.build_and(lst_block, pad, "")?
+        let pad = match direction {
+            PaddingDirection::First => builder.build_left_shift(ones, idx_in_block, "")?,
+            PaddingDirection::Last => builder.build_right_shift(ones, idx_from_end, false, "")?,
+        };
+        builder.build_and(block, pad, "")?
     };
-    builder.build_store(lst_block_addr, new_lst_block)?;
+    builder.build_store(block_addr, new_block)?;
     Ok(())
 }
 

hugr-llvm/src/extension/collections/snapshots/hugr_llvm__extension__collections__borrow_array__test__emit_clone@[email protected]

@@ -71,55 +71,54 @@ declare void @llvm.memset.p0i64.i64(i64* nocapture writeonly, i8, i64, i1 immarg
 define internal void @__barray_check_none_borrowed(i64* %0, i64 %1, i64 %2) {
   %4 = add i64 %1, %2
   %5 = udiv i64 %1, 64
-  %6 = udiv i64 %4, 64
-  %7 = getelementptr inbounds i64, i64* %0, i64 %5
-  %8 = getelementptr inbounds i64, i64* %0, i64 %6
-  %9 = load i64, i64* %7, align 4
-  %10 = load i64, i64* %8, align 4
-  %11 = urem i64 %1, 64
-  %12 = sub i64 64, %11
-  %13 = shl i64 -1, %11
-  %14 = and i64 %9, %13
-  store i64 %14, i64* %7, align 4
+  %6 = getelementptr inbounds i64, i64* %0, i64 %5
+  %7 = load i64, i64* %6, align 4
+  %8 = urem i64 %1, 64
+  %9 = sub i64 64, %8
+  %10 = shl i64 -1, %8
+  %11 = and i64 %7, %10
+  store i64 %11, i64* %6, align 4
+  %12 = udiv i64 %4, 64
+  %13 = getelementptr inbounds i64, i64* %0, i64 %12
+  %14 = load i64, i64* %13, align 4
   %15 = urem i64 %4, 64
   %16 = sub i64 64, %15
   %17 = lshr i64 -1, %16
-  %18 = and i64 %10, %17
-  store i64 %18, i64* %8, align 4
-  %19 = add i64 %1, %2
-  %20 = udiv i64 %1, 64
-  %21 = udiv i64 %19, 64
-  %22 = sub i64 %21, %20
-  %23 = add i64 %22, 1
-  %24 = alloca i64, align 8
-  store i64 0, i64* %24, align 4
-  br label %25
-
-25:                                               ; preds = %mask_block_ok, %3
-  %26 = load i64, i64* %24, align 4
-  %27 = icmp ult i64 %26, %23
-  br i1 %27, label %28, label %34
-
-28:                                               ; preds = %25
-  %29 = load i64, i64* %24, align 4
-  %30 = add i64 %29, %20
-  %31 = getelementptr inbounds i64, i64* %0, i64 %30
-  %32 = load i64, i64* %31, align 4
-  %33 = icmp eq i64 %32, 0
-  br i1 %33, label %mask_block_ok, label %mask_block_err
-
-34:                                               ; preds = %25
+  %18 = and i64 %14, %17
+  store i64 %18, i64* %13, align 4
+  %19 = udiv i64 %1, 64
+  %20 = udiv i64 %4, 64
+  %21 = sub i64 %20, %19
+  %22 = add i64 %21, 1
+  %23 = alloca i64, align 8
+  store i64 0, i64* %23, align 4
+  br label %24
+
+24:                                               ; preds = %mask_block_ok, %3
+  %25 = load i64, i64* %23, align 4
+  %26 = icmp ult i64 %25, %22
+  br i1 %26, label %27, label %33
+
+27:                                               ; preds = %24
+  %28 = load i64, i64* %23, align 4
+  %29 = add i64 %28, %19
+  %30 = getelementptr inbounds i64, i64* %0, i64 %29
+  %31 = load i64, i64* %30, align 4
+  %32 = icmp eq i64 %31, 0
+  br i1 %32, label %mask_block_ok, label %mask_block_err
+
+33:                                               ; preds = %24
   ret void
 
-mask_block_err:                                   ; preds = %28
-  %35 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([34 x i8], [34 x i8]* @prelude.panic_template, i32 0, i32 0), i32 2, i8* getelementptr inbounds ([39 x i8], [39 x i8]* @0, i32 0, i32 0))
+mask_block_err:                                   ; preds = %27
+  %34 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([34 x i8], [34 x i8]* @prelude.panic_template, i32 0, i32 0), i32 2, i8* getelementptr inbounds ([39 x i8], [39 x i8]* @0, i32 0, i32 0))
   call void @abort()
   unreachable
 
-mask_block_ok:                                    ; preds = %28
-  %36 = add i64 %29, 1
-  store i64 %36, i64* %24, align 4
-  br label %25
+mask_block_ok:                                    ; preds = %27
+  %35 = add i64 %28, 1
+  store i64 %35, i64* %23, align 4
+  br label %24
 }
 
 declare i32 @printf(i8*, ...)