Fix PassMode::Cast ABI argument lowering

Author: Kuhai9801

src/abi/pass_mode.rs

@@ -27,9 +27,12 @@ fn reg_to_abi_param(reg: Reg) -> AbiParam {
         (RegKind::Float, 4) => types::F32,
         (RegKind::Float, 8) => types::F64,
         (RegKind::Float, 16) => types::F128,
-        (RegKind::Vector { hint_vector_elem: _ }, size) => {
-            types::I8.by(u32::try_from(size).unwrap()).unwrap()
-        }
+        (
+            RegKind::Vector {
+                hint_vector_elem: _,
+            },
+            size,
+        ) => types::I8.by(u32::try_from(size).unwrap()).unwrap(),
         _ => unreachable!("{:?}", reg),
     };
     AbiParam::new(clif_ty)
@@ -55,13 +58,21 @@ fn cast_target_to_abi_params(cast: &CastTarget) -> SmallVec<[(Size, AbiParam); 2
         ];
     }
 
-    let (rest_count, rem_bytes) = if cast.rest.unit.size.bytes() == 0 {
-        (0, 0)
+    let rest_count = if cast.rest.total == Size::ZERO {
+        0
     } else {
-        (
-            cast.rest.total.bytes() / cast.rest.unit.size.bytes(),
-            cast.rest.total.bytes() % cast.rest.unit.size.bytes(),
-        )
+        assert_ne!(cast.rest.unit.size, Size::ZERO);
+        if cast.rest.total.bytes() % cast.rest.unit.size.bytes() != 0 {
+            assert_eq!(
+                cast.rest.unit.kind,
+                RegKind::Integer,
+                "only int regs can be split"
+            );
+        }
+        cast.rest
+            .total
+            .bytes()
+            .div_ceil(cast.rest.unit.size.bytes())
     };
 
     // Note: Unlike the LLVM equivalent of this code we don't have separate branches for when there
@@ -83,18 +94,39 @@ fn cast_target_to_abi_params(cast: &CastTarget) -> SmallVec<[(Size, AbiParam); 2
         res.push((offset, arg));
         offset += Size::from_bytes(arg.value_type.bytes());
     }
+    res
+}
 
-    // Append final integer
-    if rem_bytes != 0 {
-        // Only integers can be really split further.
-        assert_eq!(cast.rest.unit.kind, RegKind::Integer);
-        res.push((
-            offset,
-            reg_to_abi_param(Reg { kind: RegKind::Integer, size: Size::from_bytes(rem_bytes) }),
-        ));
-    }
+fn abi_params_size(params: &[(Size, AbiParam)]) -> u64 {
+    params
+        .iter()
+        .map(|(offset, param)| offset.bytes() + u64::from(param.value_type.bytes()))
+        .max()
+        .unwrap_or(0)
+}
 
-    res
+#[cfg(test)]
+mod tests {
+    use rustc_target::callconv::Uniform;
+
+    use super::*;
+
+    #[test]
+    fn cast_target_rounds_integer_rest_to_unit_size() {
+        let reg = Reg {
+            kind: RegKind::Integer,
+            size: Size::from_bytes(8),
+        };
+        let cast = CastTarget::from(Uniform::new(reg, Size::from_bytes(10)));
+
+        let params = cast_target_to_abi_params(&cast);
+
+        assert_eq!(params.len(), 2);
+        assert_eq!(params[0].0, Size::ZERO);
+        assert_eq!(params[0].1.value_type, types::I64);
+        assert_eq!(params[1].0, Size::from_bytes(8));
+        assert_eq!(params[1].1.value_type, types::I64);
+    }
 }
 
 impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
@@ -125,22 +157,39 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
             },
             PassMode::Cast { ref cast, pad_i32 } => {
                 assert!(!pad_i32, "padding support not yet implemented");
-                cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect()
+                cast_target_to_abi_params(cast)
+                    .into_iter()
+                    .map(|(_, param)| param)
+                    .collect()
             }
-            PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
+            PassMode::Indirect {
+                attrs,
+                meta_attrs: None,
+                on_stack,
+            } => {
                 if on_stack {
                     // Abi requires aligning struct size to pointer size
-                    let size = self.layout.size.align_to(tcx.data_layout.pointer_align().abi);
+                    let size = self
+                        .layout
+                        .size
+                        .align_to(tcx.data_layout.pointer_align().abi);
                     let size = u32::try_from(size.bytes()).unwrap();
                     smallvec![apply_attrs_to_abi_param(
                         AbiParam::special(pointer_ty(tcx), ArgumentPurpose::StructArgument(size),),
                         attrs
                     )]
                 } else {
-                    smallvec![apply_attrs_to_abi_param(AbiParam::new(pointer_ty(tcx)), attrs)]
+                    smallvec![apply_attrs_to_abi_param(
+                        AbiParam::new(pointer_ty(tcx)),
+                        attrs
+                    )]
                 }
             }
-            PassMode::Indirect { attrs, meta_attrs: Some(meta_attrs), on_stack } => {
+            PassMode::Indirect {
+                attrs,
+                meta_attrs: Some(meta_attrs),
+                on_stack,
+            } => {
                 assert!(!on_stack);
                 smallvec![
                     apply_attrs_to_abi_param(AbiParam::new(pointer_ty(tcx)), attrs),
@@ -163,7 +212,10 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 ),
                 BackendRepr::SimdVector { .. } => {
                     let vector_ty = crate::intrinsics::clif_vector_type(tcx, self.layout);
-                    (None, vec![apply_attrs_to_abi_param(AbiParam::new(vector_ty), attrs)])
+                    (
+                        None,
+                        vec![apply_attrs_to_abi_param(AbiParam::new(vector_ty), attrs)],
+                    )
                 }
                 _ => unreachable!("{:?}", self.layout.backend_repr),
             },
@@ -183,9 +235,16 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
             },
             PassMode::Cast { ref cast, .. } => (
                 None,
-                cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect(),
+                cast_target_to_abi_params(cast)
+                    .into_iter()
+                    .map(|(_, param)| param)
+                    .collect(),
             ),
-            PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
+            PassMode::Indirect {
+                attrs,
+                meta_attrs: None,
+                on_stack,
+            } => {
                 assert!(!on_stack);
                 (
                     Some(apply_attrs_to_abi_param(
@@ -195,7 +254,11 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                     vec![],
                 )
             }
-            PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+            PassMode::Indirect {
+                attrs: _,
+                meta_attrs: Some(_),
+                on_stack: _,
+            } => {
                 unreachable!("unsized return value")
             }
         }
@@ -207,12 +270,51 @@ pub(super) fn to_casted_value<'tcx>(
     arg: CValue<'tcx>,
     cast: &CastTarget,
 ) -> SmallVec<[Value; 2]> {
+    let layout = arg.layout();
     let (ptr, meta) = arg.force_stack(fx);
     assert!(meta.is_none());
-    cast_target_to_abi_params(cast)
-        .into_iter()
+    let abi_params = cast_target_to_abi_params(cast);
+    let abi_param_size = abi_params_size(&abi_params);
+    let cast_align = cast.align(fx).bytes();
+    let layout_align = layout.align.bytes();
+    let slot_align = cast_align.max(layout_align);
+    let slot_size = abi_param_size
+        .max(layout.size.bytes())
+        .next_multiple_of(slot_align);
+    let scratch = fx.create_stack_slot(
+        u32::try_from(slot_size).unwrap(),
+        u32::try_from(slot_align).unwrap(),
+    );
+    let zero = fx.bcx.ins().iconst(types::I8, 0);
+    for offset in 0..slot_size {
+        scratch
+            .offset_i64(fx, offset as i64)
+            .store(fx, zero, MemFlags::new());
+    }
+
+    let copy_bytes = abi_param_size.min(layout.size.bytes());
+    if copy_bytes != 0 {
+        let from_addr = ptr.get_addr(fx);
+        let to_addr = scratch.get_addr(fx);
+        fx.bcx.emit_small_memory_copy(
+            fx.target_config,
+            to_addr,
+            from_addr,
+            copy_bytes,
+            slot_align.try_into().unwrap_or(128),
+            layout.align.bytes().try_into().unwrap_or(128),
+            true,
+            MemFlags::new(),
+        );
+    }
+    abi_params
+        .iter()
         .map(|(offset, param)| {
-            ptr.offset_i64(fx, offset.bytes() as i64).load(fx, param.value_type, MemFlags::new())
+            scratch.offset_i64(fx, offset.bytes() as i64).load(
+                fx,
+                param.value_type,
+                MemFlags::new(),
+            )
         })
         .collect()
 }
@@ -224,14 +326,15 @@ pub(super) fn from_casted_value<'tcx>(
     cast: &CastTarget,
 ) -> CValue<'tcx> {
     let abi_params = cast_target_to_abi_params(cast);
-    let abi_param_size: u32 = abi_params.iter().map(|(_, param)| param.value_type.bytes()).sum();
-    let layout_size = u32::try_from(layout.size.bytes()).unwrap();
+    let abi_param_size = abi_params_size(&abi_params);
+    let cast_align = cast.align(fx).bytes();
+    let layout_size = layout.size.bytes();
+    let layout_align = layout.align.bytes();
+    let slot_align = cast_align.max(layout_align);
+    let slot_size = abi_param_size.max(layout_size).next_multiple_of(slot_align);
     let ptr = fx.create_stack_slot(
-        // Stack slot size may be bigger for example `[u8; 3]` which is packed into an `i32`.
-        // It may also be smaller for example when the type is a wrapper around an integer with a
-        // larger alignment than the integer.
-        std::cmp::max(abi_param_size, layout_size),
-        u32::try_from(layout.align.bytes()).unwrap(),
+        u32::try_from(slot_size).unwrap(),
+        u32::try_from(slot_align).unwrap(),
     );
     let mut block_params_iter = block_params.iter().copied();
     for (offset, _) in abi_params {
@@ -293,7 +396,10 @@ pub(super) fn cvalue_for_param<'tcx>(
         .into_iter()
         .map(|abi_param| {
             let block_param = block_params_iter.next().unwrap();
-            assert_eq!(fx.bcx.func.dfg.value_type(block_param), abi_param.value_type);
+            assert_eq!(
+                fx.bcx.func.dfg.value_type(block_param),
+                abi_param.value_type
+            );
             block_param
         })
         .collect::<SmallVec<[_; 2]>>();
@@ -321,7 +427,11 @@ pub(super) fn cvalue_for_param<'tcx>(
         PassMode::Cast { ref cast, .. } => {
             from_casted_value(fx, &block_params, arg_abi.layout, cast)
         }
-        PassMode::Indirect { attrs, meta_attrs: None, on_stack: _ } => {
+        PassMode::Indirect {
+            attrs,
+            meta_attrs: None,
+            on_stack: _,
+        } => {
             assert_eq!(block_params.len(), 1, "{:?}", block_params);
             if let Some(pointee_align) = attrs.pointee_align
                 && pointee_align < arg_abi.layout.align.abi
@@ -339,11 +449,22 @@ pub(super) fn cvalue_for_param<'tcx>(
                 CValue::by_ref(Pointer::new(block_params[0]), arg_abi.layout)
             }
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect {
+            attrs: _,
+            meta_attrs: Some(_),
+            on_stack: _,
+        } => {
             assert_eq!(block_params.len(), 2, "{:?}", block_params);
-            CValue::by_ref_unsized(Pointer::new(block_params[0]), block_params[1], arg_abi.layout)
+            CValue::by_ref_unsized(
+                Pointer::new(block_params[0]),
+                block_params[1],
+                arg_abi.layout,
+            )
         }
     };
 
-    Some(ArgValue { value, is_underaligned_pointee: false })
+    Some(ArgValue {
+        value,
+        is_underaligned_pointee: false,
+    })
 }