Auto merge of #15194 - HKalbasi:mir, r=HKalbasi

Fix layout of simd types and respect align in mir interpreter

crates/hir-def/src/data/adt.rs

@@ -147,6 +147,7 @@ fn parse_repr_tt(tt: &Subtree) -> Option<ReprOptions> {
                 }
                 "C" => ReprFlags::IS_C,
                 "transparent" => ReprFlags::IS_TRANSPARENT,
+                "simd" => ReprFlags::IS_SIMD,
                 repr => {
                     if let Some(builtin) = BuiltinInt::from_suffix(repr)
                         .map(Either::Left)

crates/hir-ty/src/consteval/tests.rs

@@ -36,6 +36,18 @@ fn check_fail(ra_fixture: &str, error: impl FnOnce(ConstEvalError) -> bool) {
 
 #[track_caller]
 fn check_number(ra_fixture: &str, answer: i128) {
+    check_answer(ra_fixture, |b| {
+        assert_eq!(
+            b,
+            &answer.to_le_bytes()[0..b.len()],
+            "Bytes differ. In decimal form: actual = {}, expected = {answer}",
+            i128::from_le_bytes(pad16(b, true))
+        );
+    });
+}
+
+#[track_caller]
+fn check_answer(ra_fixture: &str, check: impl FnOnce(&[u8])) {
     let (db, file_id) = TestDB::with_single_file(ra_fixture);
     let r = match eval_goal(&db, file_id) {
         Ok(t) => t,
@@ -47,12 +59,7 @@ fn check_number(ra_fixture: &str, answer: i128) {
     match &r.data(Interner).value {
         chalk_ir::ConstValue::Concrete(c) => match &c.interned {
             ConstScalar::Bytes(b, _) => {
-                assert_eq!(
-                    b,
-                    &answer.to_le_bytes()[0..b.len()],
-                    "Bytes differ. In decimal form: actual = {}, expected = {answer}",
-                    i128::from_le_bytes(pad16(b, true))
-                );
+                check(b);
             }
             x => panic!("Expected number but found {:?}", x),
         },
@@ -87,7 +94,7 @@ fn eval_goal(db: &TestDB, file_id: FileId) -> Result<Const, ConstEvalError> {
             }
             _ => None,
         })
-        .unwrap();
+        .expect("No const named GOAL found in the test");
     db.const_eval(const_id.into(), Substitution::empty(Interner))
 }
 
@@ -206,6 +213,30 @@ fn raw_pointer_equality() {
     );
 }
 
+#[test]
+fn alignment() {
+    check_answer(
+        r#"
+//- minicore: transmute
+use core::mem::transmute;
+const GOAL: usize = {
+    let x: i64 = 2;
+    transmute(&x)
+}
+        "#,
+        |b| assert_eq!(b[0] % 8, 0),
+    );
+    check_answer(
+        r#"
+//- minicore: transmute
+use core::mem::transmute;
+static X: i64 = 12;
+const GOAL: usize = transmute(&X);
+        "#,
+        |b| assert_eq!(b[0] % 8, 0),
+    );
+}
+
 #[test]
 fn locals() {
     check_number(

crates/hir-ty/src/layout.rs

@@ -7,7 +7,7 @@
         Abi, FieldsShape, Integer, LayoutCalculator, LayoutS, Primitive, ReprOptions, Scalar, Size,
         StructKind, TargetDataLayout, WrappingRange,
     },
-    LocalEnumVariantId, LocalFieldId,
+    LocalEnumVariantId, LocalFieldId, StructId,
 };
 use la_arena::{Idx, RawIdx};
 use stdx::never;
@@ -77,6 +77,78 @@ fn current_data_layout(&self) -> &'a TargetDataLayout {
     }
 }
 
+// FIXME: move this to the `rustc_abi`.
+fn layout_of_simd_ty(
+    db: &dyn HirDatabase,
+    id: StructId,
+    subst: &Substitution,
+    krate: CrateId,
+    dl: &TargetDataLayout,
+) -> Result<Arc<Layout>, LayoutError> {
+    let fields = db.field_types(id.into());
+
+    // Supported SIMD vectors are homogeneous ADTs with at least one field:
+    //
+    // * #[repr(simd)] struct S(T, T, T, T);
+    // * #[repr(simd)] struct S { x: T, y: T, z: T, w: T }
+    // * #[repr(simd)] struct S([T; 4])
+    //
+    // where T is a primitive scalar (integer/float/pointer).
+
+    let f0_ty = match fields.iter().next() {
+        Some(x) => x.1.clone().substitute(Interner, subst),
+        None => {
+            user_error!("simd type with zero fields");
+        }
+    };
+
+    // The element type and number of elements of the SIMD vector
+    // are obtained from:
+    //
+    // * the element type and length of the single array field, if
+    // the first field is of array type, or
+    //
+    // * the homogeneous field type and the number of fields.
+    let (e_ty, e_len, is_array) = if let TyKind::Array(e_ty, _) = f0_ty.kind(Interner) {
+        // Extract the number of elements from the layout of the array field:
+        let FieldsShape::Array { count, .. } = db.layout_of_ty(f0_ty.clone(), krate)?.fields else {
+            user_error!("Array with non array layout");
+        };
+
+        (e_ty.clone(), count, true)
+    } else {
+        // First ADT field is not an array:
+        (f0_ty, fields.iter().count() as u64, false)
+    };
+
+    // Compute the ABI of the element type:
+    let e_ly = db.layout_of_ty(e_ty, krate)?;
+    let Abi::Scalar(e_abi) = e_ly.abi else {
+        user_error!("simd type with inner non scalar type");
+    };
+
+    // Compute the size and alignment of the vector:
+    let size = e_ly.size.checked_mul(e_len, dl).ok_or(LayoutError::SizeOverflow)?;
+    let align = dl.vector_align(size);
+    let size = size.align_to(align.abi);
+
+    // Compute the placement of the vector fields:
+    let fields = if is_array {
+        FieldsShape::Arbitrary { offsets: [Size::ZERO].into(), memory_index: [0].into() }
+    } else {
+        FieldsShape::Array { stride: e_ly.size, count: e_len }
+    };
+
+    Ok(Arc::new(Layout {
+        variants: Variants::Single { index: struct_variant_idx() },
+        fields,
+        abi: Abi::Vector { element: e_abi, count: e_len },
+        largest_niche: e_ly.largest_niche,
+        size,
+        align,
+    }))
+}
+
 pub fn layout_of_ty_query(
     db: &dyn HirDatabase,
     ty: Ty,
@@ -88,7 +160,16 @@ pub fn layout_of_ty_query(
     let trait_env = Arc::new(TraitEnvironment::empty(krate));
     let ty = normalize(db, trait_env, ty.clone());
     let result = match ty.kind(Interner) {
-        TyKind::Adt(AdtId(def), subst) => return db.layout_of_adt(*def, subst.clone(), krate),
+        TyKind::Adt(AdtId(def), subst) => {
+            if let hir_def::AdtId::StructId(s) = def {
+                let data = db.struct_data(*s);
+                let repr = data.repr.unwrap_or_default();
+                if repr.simd() {
+                    return layout_of_simd_ty(db, *s, subst, krate, &target);
+                }
+            };
+            return db.layout_of_adt(*def, subst.clone(), krate);
+        }
         TyKind::Scalar(s) => match s {
             chalk_ir::Scalar::Bool => Layout::scalar(
                 dl,

crates/hir-ty/src/layout/tests.rs

@@ -270,6 +270,20 @@ impl Tr for S {
     );
 }
 
+#[test]
+fn simd_types() {
+    check_size_and_align(
+        r#"
+            #[repr(simd)]
+            struct SimdType(i64, i64);
+            struct Goal(SimdType);
+        "#,
+        "",
+        16,
+        16,
+    );
+}
+
 #[test]
 fn return_position_impl_trait() {
     size_and_align_expr! {

crates/hir-ty/src/lib.rs

@@ -180,9 +180,16 @@ fn insert(&mut self, addr: usize, x: Vec<u8>) {
     /// allocator function as `f` and it will return a mapping of old addresses to new addresses.
     fn transform_addresses(
         &self,
-        mut f: impl FnMut(&[u8]) -> Result<usize, MirEvalError>,
+        mut f: impl FnMut(&[u8], usize) -> Result<usize, MirEvalError>,
     ) -> Result<HashMap<usize, usize>, MirEvalError> {
-        self.memory.iter().map(|x| Ok((*x.0, f(x.1)?))).collect()
+        self.memory
+            .iter()
+            .map(|x| {
+                let addr = *x.0;
+                let align = if addr == 0 { 64 } else { (addr - (addr & (addr - 1))).min(64) };
+                Ok((addr, f(x.1, align)?))
+            })
+            .collect()
     }
 
     fn get<'a>(&'a self, addr: usize, size: usize) -> Option<&'a [u8]> {

crates/hir-ty/src/mir/eval.rs

@@ -226,16 +226,26 @@ fn get<'a>(&'a self, memory: &'a Evaluator<'a>) -> Result<&'a [u8]> {
     }
 }
 
+#[cfg(target_pointer_width = "64")]
+const STACK_OFFSET: usize = 1 << 60;
+#[cfg(target_pointer_width = "64")]
+const HEAP_OFFSET: usize = 1 << 59;
+
+#[cfg(target_pointer_width = "32")]
+const STACK_OFFSET: usize = 1 << 30;
+#[cfg(target_pointer_width = "32")]
+const HEAP_OFFSET: usize = 1 << 29;
+
 impl Address {
     fn from_bytes(x: &[u8]) -> Result<Self> {
         Ok(Address::from_usize(from_bytes!(usize, x)))
     }
 
     fn from_usize(x: usize) -> Self {
-        if x > usize::MAX / 2 {
-            Stack(x - usize::MAX / 2)
-        } else if x > usize::MAX / 4 {
-            Heap(x - usize::MAX / 4)
+        if x > STACK_OFFSET {
+            Stack(x - STACK_OFFSET)
+        } else if x > HEAP_OFFSET {
+            Heap(x - HEAP_OFFSET)
         } else {
             Invalid(x)
         }
@@ -247,8 +257,8 @@ fn to_bytes(&self) -> Vec<u8> {
 
     fn to_usize(&self) -> usize {
         let as_num = match self {
-            Stack(x) => *x + usize::MAX / 2,
-            Heap(x) => *x + usize::MAX / 4,
+            Stack(x) => *x + STACK_OFFSET,
+            Heap(x) => *x + HEAP_OFFSET,
             Invalid(x) => *x,
         };
         as_num
@@ -721,8 +731,14 @@ fn interpret_mir(
                 .locals
                 .iter()
                 .map(|(id, x)| {
-                    let size =
-                        self.size_of_sized(&x.ty, &locals, "no unsized local in extending stack")?;
+                    let (size, align) = self.size_align_of_sized(
+                        &x.ty,
+                        &locals,
+                        "no unsized local in extending stack",
+                    )?;
+                    while stack_ptr % align != 0 {
+                        stack_ptr += 1;
+                    }
                     let my_ptr = stack_ptr;
                     stack_ptr += size;
                     Ok((id, Interval { addr: Stack(my_ptr), size }))
@@ -1469,8 +1485,8 @@ fn allocate_const_in_heap(
         Ok(match &c.interned {
             ConstScalar::Bytes(v, memory_map) => {
                 let mut v: Cow<'_, [u8]> = Cow::Borrowed(v);
-                let patch_map = memory_map.transform_addresses(|b| {
-                    let addr = self.heap_allocate(b.len(), 1); // FIXME: align is wrong
+                let patch_map = memory_map.transform_addresses(|b, align| {
+                    let addr = self.heap_allocate(b.len(), align);
                     self.write_memory(addr, b)?;
                     Ok(addr.to_usize())
                 })?;
@@ -1574,7 +1590,24 @@ fn size_of_sized(&self, ty: &Ty, locals: &Locals<'_>, what: &'static str) -> Res
         }
     }
 
-    fn heap_allocate(&mut self, size: usize, _align: usize) -> Address {
+    /// A version of `self.size_align_of` which returns error if the type is unsized. `what` argument should
+    /// be something that complete this: `error: type {ty} was unsized. {what} should be sized`
+    fn size_align_of_sized(
+        &self,
+        ty: &Ty,
+        locals: &Locals<'_>,
+        what: &'static str,
+    ) -> Result<(usize, usize)> {
+        match self.size_align_of(ty, locals)? {
+            Some(x) => Ok(x),
+            None => Err(MirEvalError::TypeIsUnsized(ty.clone(), what)),
+        }
+    }
+
+    fn heap_allocate(&mut self, size: usize, align: usize) -> Address {
+        while self.heap.len() % align != 0 {
+            self.heap.push(0);
+        }
         let pos = self.heap.len();
         self.heap.extend(iter::repeat(0).take(size));
         Address::Heap(pos)