x86_64 sysv64 vararg: support structs

folkertdev · folkertdev · commit dcdc64554fb8 · 2025-05-26T13:34:14.000+02:00
diff --git a/compiler/rustc_codegen_llvm/src/va_arg.rs b/compiler/rustc_codegen_llvm/src/va_arg.rs
@@ -1,7 +1,10 @@
-use rustc_abi::{Align, BackendRepr, Endian, HasDataLayout, Primitive, Size};
+use rustc_abi::{Align, BackendRepr, Endian, HasDataLayout, Primitive, Size, TyAndLayout};
+use rustc_codegen_ssa::MemFlags;
 use rustc_codegen_ssa::common::IntPredicate;
 use rustc_codegen_ssa::mir::operand::OperandRef;
-use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods, ConstCodegenMethods};
+use rustc_codegen_ssa::traits::{
+    BaseTypeCodegenMethods, BuilderMethods, ConstCodegenMethods, LayoutTypeCodegenMethods,
+};
 use rustc_middle::ty::Ty;
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
 
@@ -300,11 +303,6 @@ fn emit_x86_64_sysv64_va_arg<'ll, 'tcx>(
     // } va_list[1];
     let va_list_addr = list.immediate();
 
-    let unsigned_int_offset = 4;
-    let ptr_offset = 8;
-    let gp_offset_ptr = va_list_addr;
-    let fp_offset_ptr = bx.inbounds_ptradd(va_list_addr, bx.cx.const_usize(unsigned_int_offset));
-
     let layout = bx.cx.layout_of(target_ty);
 
     // AMD64-ABI 3.5.7p5: Step 1. Determine whether type may be passed
@@ -317,37 +315,48 @@ fn emit_x86_64_sysv64_va_arg<'ll, 'tcx>(
     let mut num_gp_registers = 0;
     let mut num_fp_registers = 0;
 
+    let mut registers_for_primitive = |p| match p {
+        Primitive::Int(integer, _is_signed) => {
+            num_gp_registers += integer.size().bytes().div_ceil(8) as u32;
+        }
+        Primitive::Float(float) => {
+            num_fp_registers += float.size().bytes().div_ceil(16) as u32;
+        }
+        Primitive::Pointer(_) => {
+            num_gp_registers += 1;
+        }
+    };
+
     match layout.layout.backend_repr() {
-        BackendRepr::Scalar(scalar) => match scalar.primitive() {
-            Primitive::Int(integer, _is_signed) => {
-                num_gp_registers += integer.size().bytes().div_ceil(8) as u32;
-            }
-            Primitive::Float(float) => {
-                num_fp_registers += float.size().bytes().div_ceil(16) as u32;
-            }
-            Primitive::Pointer(_) => {
-                num_gp_registers += 1;
-            }
-        },
-        BackendRepr::ScalarPair(..)
-        | BackendRepr::SimdVector { .. }
-        | BackendRepr::Memory { .. } => {
+        BackendRepr::Scalar(scalar) => {
+            registers_for_primitive(scalar.primitive());
+        }
+        BackendRepr::ScalarPair(scalar1, scalar2) => {
+            registers_for_primitive(scalar1.primitive());
+            registers_for_primitive(scalar2.primitive());
+        }
+        BackendRepr::SimdVector { .. } => {
             // Because no instance of VaArgSafe uses a non-scalar `BackendRepr`.
             unreachable!(
                 "No x86-64 SysV va_arg implementation for {:?}",
                 layout.layout.backend_repr()
             )
         }
+        BackendRepr::Memory { .. } => {
+            let mem_addr = x86_64_sysv64_va_arg_from_memory(bx, va_list_addr, layout);
+            return bx.load(layout.llvm_type(bx), mem_addr, layout.align.abi);
+        }
     };
 
-    if num_gp_registers == 0 && num_fp_registers == 0 {
-        unreachable!("VaArgSafe is not implemented for ZSTs")
-    }
-
     // AMD64-ABI 3.5.7p5: Step 3. Verify whether arguments fit into
     // registers. In the case: l->gp_offset > 48 - num_gp * 8 or
     // l->fp_offset > 176 - num_fp * 16 go to step 7.
 
+    let unsigned_int_offset = 4;
+    let ptr_offset = 8;
+    let gp_offset_ptr = va_list_addr;
+    let fp_offset_ptr = bx.inbounds_ptradd(va_list_addr, bx.cx.const_usize(unsigned_int_offset));
+
     let gp_offset_v = bx.load(bx.type_i32(), gp_offset_ptr, Align::from_bytes(8).unwrap());
     let fp_offset_v = bx.load(bx.type_i32(), fp_offset_ptr, Align::from_bytes(4).unwrap());
 
@@ -388,14 +397,85 @@ fn emit_x86_64_sysv64_va_arg<'ll, 'tcx>(
         bx.inbounds_ptradd(va_list_addr, bx.cx.const_usize(2 * unsigned_int_offset + ptr_offset));
     let reg_save_area_v = bx.load(bx.type_ptr(), reg_save_area_ptr, dl.pointer_align.abi);
 
-    let reg_addr = if num_gp_registers > 0 && num_fp_registers > 0 {
-        unreachable!("instances of VaArgSafe cannot use both int and sse registers");
-    } else if num_gp_registers > 0 || num_fp_registers == 1 {
-        let gp_or_fp_offset = if num_gp_registers > 0 { gp_offset_v } else { fp_offset_v };
-        bx.gep(bx.type_i8(), reg_save_area_v, &[gp_or_fp_offset])
-    } else {
-        // assert_eq!(num_sse_registers, 2);
-        unreachable!("all instances of VaArgSafe have an alignment <= 8");
+    let reg_addr = match layout.layout.backend_repr() {
+        BackendRepr::Scalar(scalar) => match scalar.primitive() {
+            Primitive::Int(_, _) | Primitive::Pointer(_) => {
+                let reg_addr = bx.gep(bx.type_i8(), reg_save_area_v, &[gp_offset_v]);
+
+                // Copy into a temporary if the type is more aligned than the register save area.
+                copy_to_temporary_if_more_aligned(bx, reg_addr, layout)
+            }
+            Primitive::Float(_) => bx.gep(bx.type_i8(), reg_save_area_v, &[fp_offset_v]),
+        },
+        BackendRepr::ScalarPair(scalar1, scalar2) => {
+            let ty_lo = bx.cx().scalar_pair_element_backend_type(layout, 0, false);
+            let ty_hi = bx.cx().scalar_pair_element_backend_type(layout, 1, false);
+
+            let align_lo = layout.field(bx.cx, 0).layout.align().abi;
+            let align_hi = layout.field(bx.cx, 1).layout.align().abi;
+
+            match (scalar1.primitive(), scalar2.primitive()) {
+                (Primitive::Float(_), Primitive::Float(_)) => {
+                    // SSE registers are spaced 16 bytes apart in the register save
+                    // area, we need to collect the two eightbytes together.
+                    // The ABI isn't explicit about this, but it seems reasonable
+                    // to assume that the slots are 16-byte aligned, since the stack is
+                    // naturally 16-byte aligned and the prologue is expected to store
+                    // all the SSE registers to the RSA.
+                    let reg_lo_addr = bx.gep(bx.type_i8(), reg_save_area_v, &[fp_offset_v]);
+                    let reg_hi_addr = bx.gep(bx.type_i8(), reg_lo_addr, &[bx.const_i32(16)]);
+
+                    let tmp = bx.alloca(layout.layout.size(), layout.layout.align().abi);
+                    let align = layout.layout.align().abi;
+
+                    let reg_lo = bx.load(ty_lo, reg_lo_addr, align_lo);
+                    let reg_hi = bx.load(ty_hi, reg_hi_addr, align_hi);
+
+                    let field0 = tmp;
+                    let field1 = bx.gep(bx.type_i8(), tmp, &[bx.const_i32(8)]);
+
+                    bx.store(reg_lo, field0, align);
+                    bx.store(reg_hi, field1, align);
+
+                    tmp
+                }
+                (Primitive::Float(_), _) | (_, Primitive::Float(_)) => {
+                    let gp_addr = bx.gep(bx.type_i8(), reg_save_area_v, &[gp_offset_v]);
+                    let fp_addr = bx.gep(bx.type_i8(), reg_save_area_v, &[fp_offset_v]);
+
+                    let (reg_lo_addr, reg_hi_addr) = match scalar1.primitive() {
+                        Primitive::Float(_) => (fp_addr, gp_addr),
+                        Primitive::Int(_, _) | Primitive::Pointer(_) => (gp_addr, fp_addr),
+                    };
+
+                    let tmp = bx.alloca(layout.layout.size(), layout.layout.align().abi);
+
+                    let reg_lo = bx.load(ty_lo, reg_lo_addr, align_lo);
+                    let reg_hi = bx.load(ty_hi, reg_hi_addr, align_hi);
+
+                    let field0 = tmp;
+                    let field1 = bx.gep(bx.type_i8(), tmp, &[bx.const_i32(8)]);
+
+                    bx.store(reg_lo, field0, layout.layout.align().abi);
+                    bx.store(reg_hi, field1, layout.layout.align().abi);
+
+                    tmp
+                }
+                (_, _) => {
+                    // Two integer/pointer values are just contiguous in memory.
+                    let reg_addr = bx.gep(bx.type_i8(), reg_save_area_v, &[gp_offset_v]);
+
+                    // Copy into a temporary if the type is more aligned than the register save area.
+                    copy_to_temporary_if_more_aligned(bx, reg_addr, layout)
+                }
+            }
+        }
+        BackendRepr::SimdVector { .. } => {
+            unreachable!("panics in the previous match on `backend_repr`")
+        }
+        BackendRepr::Memory { .. } => {
+            unreachable!("early returns in the previous match on `backend_repr`")
+        }
     };
 
     // AMD64-ABI 3.5.7p5: Step 5. Set:
@@ -416,9 +496,47 @@ fn emit_x86_64_sysv64_va_arg<'ll, 'tcx>(
     bx.br(end);
 
     bx.switch_to_block(in_mem);
+    let mem_addr = x86_64_sysv64_va_arg_from_memory(bx, va_list_addr, layout);
+    bx.br(end);
+
+    bx.switch_to_block(end);
+
+    let val_type = layout.llvm_type(bx);
+    let val_addr = bx.phi(bx.type_ptr(), &[reg_addr, mem_addr], &[in_reg, in_mem]);
 
-    let overflow_arg_area_ptr =
-        bx.inbounds_ptradd(va_list_addr, bx.cx.const_usize(2 * unsigned_int_offset));
+    bx.load(val_type, val_addr, layout.align.abi)
+}
+
+/// Copy into a temporary if the type is more aligned than the register save area.
+fn copy_to_temporary_if_more_aligned<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    reg_addr: &'ll Value,
+    layout: TyAndLayout<'tcx, Ty<'tcx>>,
+) -> &'ll Value {
+    if layout.layout.align.abi.bytes() > 8 {
+        let tmp = bx.alloca(layout.layout.size(), layout.layout.align().abi);
+        bx.memcpy(
+            tmp,
+            layout.layout.align.abi,
+            reg_addr,
+            Align::from_bytes(8).unwrap(),
+            bx.const_u32(layout.layout.size().bytes() as u32),
+            MemFlags::empty(),
+        );
+        tmp
+    } else {
+        reg_addr
+    }
+}
+
+fn x86_64_sysv64_va_arg_from_memory<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    va_list_addr: &'ll Value,
+    layout: TyAndLayout<'tcx, Ty<'tcx>>,
+) -> &'ll Value {
+    let dl = bx.cx.data_layout();
+
+    let overflow_arg_area_ptr = bx.inbounds_ptradd(va_list_addr, bx.const_usize(8));
 
     let overflow_arg_area_v = bx.load(bx.type_ptr(), overflow_arg_area_ptr, dl.pointer_align.abi);
     // AMD64-ABI 3.5.7p5: Step 7. Align l->overflow_arg_area upwards to a 16
@@ -441,14 +559,7 @@ fn emit_x86_64_sysv64_va_arg<'ll, 'tcx>(
     let overflow_arg_area = bx.gep(bx.type_i8(), overflow_arg_area_v, &[offset]);
     bx.store(overflow_arg_area, overflow_arg_area_ptr, dl.pointer_align.abi);
 
-    bx.br(end);
-
-    bx.switch_to_block(end);
-
-    let val_type = layout.llvm_type(bx);
-    let val_addr = bx.phi(bx.type_ptr(), &[reg_addr, mem_addr], &[in_reg, in_mem]);
-
-    bx.load(val_type, val_addr, layout.align.abi)
+    mem_addr
 }
 
 fn emit_xtensa_va_arg<'ll, 'tcx>(
