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

Pass `TraitEnvironment` into `layout_ty` and `const_eval`

We need to do either this or get rid of trait environment in `normalize_ty`. Let's go with this for now.

crates/hir-ty/src/consteval.rs

@@ -16,7 +16,8 @@ use triomphe::Arc;
 use crate::{
     db::HirDatabase, infer::InferenceContext, lower::ParamLoweringMode,
     mir::monomorphize_mir_body_bad, to_placeholder_idx, utils::Generics, Const, ConstData,
-    ConstScalar, ConstValue, GenericArg, Interner, MemoryMap, Substitution, Ty, TyBuilder,
+    ConstScalar, ConstValue, GenericArg, Interner, MemoryMap, Substitution, TraitEnvironment, Ty,
+    TyBuilder,
 };
 
 use super::mir::{interpret_mir, lower_to_mir, pad16, MirEvalError, MirLowerError};
@@ -135,7 +136,7 @@ pub fn intern_const_ref(
     ty: Ty,
     krate: CrateId,
 ) -> Const {
-    let layout = db.layout_of_ty(ty.clone(), krate);
+    let layout = db.layout_of_ty(ty.clone(), Arc::new(TraitEnvironment::empty(krate)));
     let bytes = match value {
         LiteralConstRef::Int(i) => {
             // FIXME: We should handle failure of layout better.
@@ -173,7 +174,7 @@ pub fn try_const_usize(db: &dyn HirDatabase, c: &Const) -> Option<u128> {
         chalk_ir::ConstValue::Concrete(c) => match &c.interned {
             ConstScalar::Bytes(it, _) => Some(u128::from_le_bytes(pad16(&it, false))),
             ConstScalar::UnevaluatedConst(c, subst) => {
-                let ec = db.const_eval(*c, subst.clone()).ok()?;
+                let ec = db.const_eval(*c, subst.clone(), None).ok()?;
                 try_const_usize(db, &ec)
             }
             _ => None,
@@ -186,6 +187,7 @@ pub(crate) fn const_eval_recover(
     _: &[String],
     _: &GeneralConstId,
     _: &Substitution,
+    _: &Option<Arc<TraitEnvironment>>,
 ) -> Result<Const, ConstEvalError> {
     Err(ConstEvalError::MirLowerError(MirLowerError::Loop))
 }
@@ -210,6 +212,7 @@ pub(crate) fn const_eval_query(
     db: &dyn HirDatabase,
     def: GeneralConstId,
     subst: Substitution,
+    trait_env: Option<Arc<TraitEnvironment>>,
 ) -> Result<Const, ConstEvalError> {
     let body = match def {
         GeneralConstId::ConstId(c) => {
@@ -228,7 +231,7 @@ pub(crate) fn const_eval_query(
         }
         GeneralConstId::InTypeConstId(c) => db.mir_body(c.into())?,
     };
-    let c = interpret_mir(db, body, false).0?;
+    let c = interpret_mir(db, body, false, trait_env).0?;
     Ok(c)
 }
 
@@ -241,7 +244,7 @@ pub(crate) fn const_eval_static_query(
         Substitution::empty(Interner),
         db.trait_environment_for_body(def.into()),
     )?;
-    let c = interpret_mir(db, body, false).0?;
+    let c = interpret_mir(db, body, false, None).0?;
     Ok(c)
 }
 
@@ -268,7 +271,7 @@ pub(crate) fn const_eval_discriminant_variant(
         Substitution::empty(Interner),
         db.trait_environment_for_body(def),
     )?;
-    let c = interpret_mir(db, mir_body, false).0?;
+    let c = interpret_mir(db, mir_body, false, None).0?;
     let c = try_const_usize(db, &c).unwrap() as i128;
     Ok(c)
 }
@@ -293,7 +296,7 @@ pub(crate) fn eval_to_const(
     }
     let infer = ctx.clone().resolve_all();
     if let Ok(mir_body) = lower_to_mir(ctx.db, ctx.owner, &ctx.body, &infer, expr) {
-        if let Ok(result) = interpret_mir(db, Arc::new(mir_body), true).0 {
+        if let Ok(result) = interpret_mir(db, Arc::new(mir_body), true, None).0 {
             return result;
         }
     }

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

@@ -114,7 +114,7 @@ fn eval_goal(db: &TestDB, file_id: FileId) -> Result<Const, ConstEvalError> {
             _ => None,
         })
         .expect("No const named GOAL found in the test");
-    db.const_eval(const_id.into(), Substitution::empty(Interner))
+    db.const_eval(const_id.into(), Substitution::empty(Interner), None)
 }
 
 #[test]
@@ -2492,6 +2492,28 @@ fn const_trait_assoc() {
     "#,
         5,
     );
+    check_number(
+        r#"
+    //- minicore: size_of
+    //- /a/lib.rs crate:a
+    use core::mem::size_of;
+    pub struct S<T>(T);
+    impl<T> S<T> {
+        pub const X: usize = core::mem::size_of::<T>();
+    }
+    //- /main.rs crate:main deps:a
+    use a::{S};
+    trait Tr {
+        type Ty;
+    }
+    impl Tr for i32 {
+        type Ty = u64;
+    }
+    struct K<T: Tr>(<T as Tr>::Ty);
+    const GOAL: usize = S::<K<i32>>::X;
+    "#,
+        8,
+    );
     check_number(
         r#"
     struct S<T>(*mut T);

crates/hir-ty/src/db.rs

@@ -77,8 +77,12 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
 
     #[salsa::invoke(crate::consteval::const_eval_query)]
     #[salsa::cycle(crate::consteval::const_eval_recover)]
-    fn const_eval(&self, def: GeneralConstId, subst: Substitution)
-        -> Result<Const, ConstEvalError>;
+    fn const_eval(
+        &self,
+        def: GeneralConstId,
+        subst: Substitution,
+        trait_env: Option<Arc<crate::TraitEnvironment>>,
+    ) -> Result<Const, ConstEvalError>;
 
     #[salsa::invoke(crate::consteval::const_eval_static_query)]
     #[salsa::cycle(crate::consteval::const_eval_static_recover)]
@@ -100,12 +104,16 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
         &self,
         def: AdtId,
         subst: Substitution,
-        krate: CrateId,
+        env: Arc<crate::TraitEnvironment>,
     ) -> Result<Arc<Layout>, LayoutError>;
 
     #[salsa::invoke(crate::layout::layout_of_ty_query)]
     #[salsa::cycle(crate::layout::layout_of_ty_recover)]
-    fn layout_of_ty(&self, ty: Ty, krate: CrateId) -> Result<Arc<Layout>, LayoutError>;
+    fn layout_of_ty(
+        &self,
+        ty: Ty,
+        env: Arc<crate::TraitEnvironment>,
+    ) -> Result<Arc<Layout>, LayoutError>;
 
     #[salsa::invoke(crate::layout::target_data_layout_query)]
     fn target_data_layout(&self, krate: CrateId) -> Option<Arc<TargetDataLayout>>;

crates/hir-ty/src/display.rs

@@ -29,6 +29,7 @@ use itertools::Itertools;
 use la_arena::ArenaMap;
 use smallvec::SmallVec;
 use stdx::never;
+use triomphe::Arc;
 
 use crate::{
     consteval::try_const_usize,
@@ -43,7 +44,7 @@ use crate::{
     AdtId, AliasEq, AliasTy, Binders, CallableDefId, CallableSig, Const, ConstScalar, ConstValue,
     DomainGoal, GenericArg, ImplTraitId, Interner, Lifetime, LifetimeData, LifetimeOutlives,
     MemoryMap, Mutability, OpaqueTy, ProjectionTy, ProjectionTyExt, QuantifiedWhereClause, Scalar,
-    Substitution, TraitRef, TraitRefExt, Ty, TyExt, WhereClause,
+    Substitution, TraitEnvironment, TraitRef, TraitRefExt, Ty, TyExt, WhereClause,
 };
 
 pub trait HirWrite: fmt::Write {
@@ -454,7 +455,9 @@ fn render_const_scalar(
 ) -> Result<(), HirDisplayError> {
     // FIXME: We need to get krate from the final callers of the hir display
     // infrastructure and have it here as a field on `f`.
-    let krate = *f.db.crate_graph().crates_in_topological_order().last().unwrap();
+    let trait_env = Arc::new(TraitEnvironment::empty(
+        *f.db.crate_graph().crates_in_topological_order().last().unwrap(),
+    ));
     match ty.kind(Interner) {
         TyKind::Scalar(s) => match s {
             Scalar::Bool => write!(f, "{}", if b[0] == 0 { false } else { true }),
@@ -497,7 +500,7 @@ fn render_const_scalar(
             TyKind::Slice(ty) => {
                 let addr = usize::from_le_bytes(b[0..b.len() / 2].try_into().unwrap());
                 let count = usize::from_le_bytes(b[b.len() / 2..].try_into().unwrap());
-                let Ok(layout) = f.db.layout_of_ty(ty.clone(), krate) else {
+                let Ok(layout) = f.db.layout_of_ty(ty.clone(), trait_env) else {
                     return f.write_str("<layout-error>");
                 };
                 let size_one = layout.size.bytes_usize();
@@ -523,7 +526,7 @@ fn render_const_scalar(
                 let Ok(t) = memory_map.vtable.ty(ty_id) else {
                     return f.write_str("<ty-missing-in-vtable-map>");
                 };
-                let Ok(layout) = f.db.layout_of_ty(t.clone(), krate) else {
+                let Ok(layout) = f.db.layout_of_ty(t.clone(), trait_env) else {
                     return f.write_str("<layout-error>");
                 };
                 let size = layout.size.bytes_usize();
@@ -555,7 +558,7 @@ fn render_const_scalar(
                         return f.write_str("<layout-error>");
                     }
                 });
-                let Ok(layout) = f.db.layout_of_ty(t.clone(), krate) else {
+                let Ok(layout) = f.db.layout_of_ty(t.clone(), trait_env) else {
                     return f.write_str("<layout-error>");
                 };
                 let size = layout.size.bytes_usize();
@@ -567,7 +570,7 @@ fn render_const_scalar(
             }
         },
         TyKind::Tuple(_, subst) => {
-            let Ok(layout) = f.db.layout_of_ty(ty.clone(), krate) else {
+            let Ok(layout) = f.db.layout_of_ty(ty.clone(), trait_env.clone()) else {
                 return f.write_str("<layout-error>");
             };
             f.write_str("(")?;
@@ -580,7 +583,7 @@ fn render_const_scalar(
                 }
                 let ty = ty.assert_ty_ref(Interner); // Tuple only has type argument
                 let offset = layout.fields.offset(id).bytes_usize();
-                let Ok(layout) = f.db.layout_of_ty(ty.clone(), krate) else {
+                let Ok(layout) = f.db.layout_of_ty(ty.clone(), trait_env.clone()) else {
                     f.write_str("<layout-error>")?;
                     continue;
                 };
@@ -590,7 +593,7 @@ fn render_const_scalar(
             f.write_str(")")
         }
         TyKind::Adt(adt, subst) => {
-            let Ok(layout) = f.db.layout_of_adt(adt.0, subst.clone(), krate) else {
+            let Ok(layout) = f.db.layout_of_adt(adt.0, subst.clone(), trait_env.clone()) else {
                 return f.write_str("<layout-error>");
             };
             match adt.0 {
@@ -602,7 +605,7 @@ fn render_const_scalar(
                         &data.variant_data,
                         f,
                         &field_types,
-                        adt.0.module(f.db.upcast()).krate(),
+                        f.db.trait_environment(adt.0.into()),
                         &layout,
                         subst,
                         b,
@@ -614,7 +617,7 @@ fn render_const_scalar(
                 }
                 hir_def::AdtId::EnumId(e) => {
                     let Some((var_id, var_layout)) =
-                        detect_variant_from_bytes(&layout, f.db, krate, b, e)
+                        detect_variant_from_bytes(&layout, f.db, trait_env.clone(), b, e)
                     else {
                         return f.write_str("<failed-to-detect-variant>");
                     };
@@ -626,7 +629,7 @@ fn render_const_scalar(
                         &data.variant_data,
                         f,
                         &field_types,
-                        adt.0.module(f.db.upcast()).krate(),
+                        f.db.trait_environment(adt.0.into()),
                         &var_layout,
                         subst,
                         b,
@@ -645,7 +648,7 @@ fn render_const_scalar(
             let Some(len) = try_const_usize(f.db, len) else {
                 return f.write_str("<unknown-array-len>");
             };
-            let Ok(layout) = f.db.layout_of_ty(ty.clone(), krate) else {
+            let Ok(layout) = f.db.layout_of_ty(ty.clone(), trait_env) else {
                 return f.write_str("<layout-error>");
             };
             let size_one = layout.size.bytes_usize();
@@ -684,7 +687,7 @@ fn render_variant_after_name(
     data: &VariantData,
     f: &mut HirFormatter<'_>,
     field_types: &ArenaMap<LocalFieldId, Binders<Ty>>,
-    krate: CrateId,
+    trait_env: Arc<TraitEnvironment>,
     layout: &Layout,
     subst: &Substitution,
     b: &[u8],
@@ -695,7 +698,7 @@ fn render_variant_after_name(
             let render_field = |f: &mut HirFormatter<'_>, id: LocalFieldId| {
                 let offset = layout.fields.offset(u32::from(id.into_raw()) as usize).bytes_usize();
                 let ty = field_types[id].clone().substitute(Interner, subst);
-                let Ok(layout) = f.db.layout_of_ty(ty.clone(), krate) else {
+                let Ok(layout) = f.db.layout_of_ty(ty.clone(), trait_env.clone()) else {
                     return f.write_str("<layout-error>");
                 };
                 let size = layout.size.bytes_usize();

crates/hir-ty/src/infer/unify.rs

@@ -252,7 +252,8 @@ impl<'a> InferenceTable<'a> {
                             // and registering an obligation. But it needs chalk support, so we handle the most basic
                             // case (a non associated const without generic parameters) manually.
                             if subst.len(Interner) == 0 {
-                                if let Ok(eval) = self.db.const_eval((*c_id).into(), subst.clone())
+                                if let Ok(eval) =
+                                    self.db.const_eval((*c_id).into(), subst.clone(), None)
                                 {
                                     eval
                                 } else {
@@ -785,7 +786,7 @@ impl<'a> InferenceTable<'a> {
                 crate::ConstScalar::Unknown => self.new_const_var(data.ty.clone()),
                 // try to evaluate unevaluated const. Replace with new var if const eval failed.
                 crate::ConstScalar::UnevaluatedConst(id, subst) => {
-                    if let Ok(eval) = self.db.const_eval(*id, subst.clone()) {
+                    if let Ok(eval) = self.db.const_eval(*id, subst.clone(), None) {
                         eval
                     } else {
                         self.new_const_var(data.ty.clone())