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

Implement recursion in mir interpreter without recursion

This enables interpreting functions with deep stack + profiling. I also applied some changes to make it faster based on the profiling result.

crates/hir-def/src/db.rs

@@ -18,7 +18,7 @@
     generics::GenericParams,
     import_map::ImportMap,
     item_tree::{AttrOwner, ItemTree},
-    lang_item::{LangItem, LangItemTarget, LangItems},
+    lang_item::{self, LangItem, LangItemTarget, LangItems},
     nameres::{diagnostics::DefDiagnostic, DefMap},
     visibility::{self, Visibility},
     AttrDefId, BlockId, BlockLoc, ConstBlockId, ConstBlockLoc, ConstId, ConstLoc, DefWithBodyId,
@@ -204,6 +204,9 @@ fn fields_attrs_source_map(
     #[salsa::invoke(AttrsWithOwner::attrs_query)]
     fn attrs(&self, def: AttrDefId) -> Attrs;
 
+    #[salsa::invoke(lang_item::lang_attr_query)]
+    fn lang_attr(&self, def: AttrDefId) -> Option<LangItem>;
+
     #[salsa::transparent]
     #[salsa::invoke(AttrsWithOwner::attrs_with_owner)]
     fn attrs_with_owner(&self, def: AttrDefId) -> AttrsWithOwner;

crates/hir-def/src/lang_item.rs

@@ -180,15 +180,15 @@ fn collect_lang_item<T>(
         T: Into<AttrDefId> + Copy,
     {
         let _p = profile::span("collect_lang_item");
-        if let Some(lang_item) = lang_attr(db, item) {
+        if let Some(lang_item) = db.lang_attr(item.into()) {
             self.items.entry(lang_item).or_insert_with(|| constructor(item));
         }
     }
 }
 
-pub fn lang_attr(db: &dyn DefDatabase, item: impl Into<AttrDefId> + Copy) -> Option<LangItem> {
-    let attrs = db.attrs(item.into());
-    attrs.by_key("lang").string_value().cloned().and_then(|it| LangItem::from_str(&it))
+pub(crate) fn lang_attr_query(db: &dyn DefDatabase, item: AttrDefId) -> Option<LangItem> {
+    let attrs = db.attrs(item);
+    attrs.by_key("lang").string_value().and_then(|it| LangItem::from_str(&it))
 }
 
 pub enum GenericRequirement {

crates/hir-ty/src/chalk_db.rs

@@ -11,7 +11,7 @@
 use base_db::CrateId;
 use hir_def::{
     hir::Movability,
-    lang_item::{lang_attr, LangItem, LangItemTarget},
+    lang_item::{LangItem, LangItemTarget},
     AssocItemId, BlockId, GenericDefId, HasModule, ItemContainerId, Lookup, TypeAliasId,
 };
 use hir_expand::name::name;
@@ -565,7 +565,7 @@ pub(crate) fn trait_datum_query(
     let where_clauses = convert_where_clauses(db, trait_.into(), &bound_vars);
     let associated_ty_ids = trait_data.associated_types().map(to_assoc_type_id).collect();
     let trait_datum_bound = rust_ir::TraitDatumBound { where_clauses };
-    let well_known = lang_attr(db.upcast(), trait_).and_then(well_known_trait_from_lang_item);
+    let well_known = db.lang_attr(trait_.into()).and_then(well_known_trait_from_lang_item);
     let trait_datum = TraitDatum {
         id: trait_id,
         binders: make_binders(db, &generic_params, trait_datum_bound),

crates/hir-ty/src/consteval.rs

@@ -228,7 +228,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).0?;
     Ok(c)
 }
 
@@ -241,7 +241,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).0?;
     Ok(c)
 }
 
@@ -268,7 +268,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).0?;
     let c = try_const_usize(db, &c).unwrap() as i128;
     Ok(c)
 }
@@ -293,7 +293,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, &mir_body, true).0 {
+        if let Ok(result) = interpret_mir(db, Arc::new(mir_body), true).0 {
             return result;
         }
     }

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

@@ -16,7 +16,7 @@
 
 fn simplify(e: ConstEvalError) -> ConstEvalError {
     match e {
-        ConstEvalError::MirEvalError(MirEvalError::InFunction(_, e, _, _)) => {
+        ConstEvalError::MirEvalError(MirEvalError::InFunction(e, _)) => {
             simplify(ConstEvalError::MirEvalError(*e))
         }
         _ => e,
@@ -2471,7 +2471,7 @@ const fn f(x: i32) -> i32 {
     }
     const GOAL: i32 = f(0);
     "#,
-        |e| e == ConstEvalError::MirEvalError(MirEvalError::StackOverflow),
+        |e| e == ConstEvalError::MirEvalError(MirEvalError::ExecutionLimitExceeded),
     );
     // Reasonable code should still work
     check_number(

crates/hir-ty/src/db.rs

@@ -110,6 +110,14 @@ fn layout_of_adt(
     #[salsa::invoke(crate::layout::target_data_layout_query)]
     fn target_data_layout(&self, krate: CrateId) -> Option<Arc<TargetDataLayout>>;
 
+    #[salsa::invoke(crate::method_resolution::lookup_impl_method_query)]
+    fn lookup_impl_method(
+        &self,
+        env: Arc<crate::TraitEnvironment>,
+        func: FunctionId,
+        fn_subst: Substitution,
+    ) -> (FunctionId, Substitution);
+
     #[salsa::invoke(crate::lower::callable_item_sig)]
     fn callable_item_signature(&self, def: CallableDefId) -> PolyFnSig;
 

crates/hir-ty/src/lower.rs

@@ -23,7 +23,7 @@
     generics::{
         TypeOrConstParamData, TypeParamProvenance, WherePredicate, WherePredicateTypeTarget,
     },
-    lang_item::{lang_attr, LangItem},
+    lang_item::LangItem,
     nameres::MacroSubNs,
     path::{GenericArg, GenericArgs, ModPath, Path, PathKind, PathSegment, PathSegments},
     resolver::{HasResolver, Resolver, TypeNs},
@@ -1012,7 +1012,7 @@ pub(crate) fn lower_type_bound(
                         //   (So ideally, we'd only ignore `~const Drop` here)
                         // - `Destruct` impls are built-in in 1.62 (current nightly as of 08-04-2022), so until
                         //   the builtin impls are supported by Chalk, we ignore them here.
-                        if let Some(lang) = lang_attr(self.db.upcast(), tr.hir_trait_id()) {
+                        if let Some(lang) = self.db.lang_attr(tr.hir_trait_id().into()) {
                             if matches!(lang, LangItem::Drop | LangItem::Destruct) {
                                 return false;
                             }

crates/hir-ty/src/method_resolution.rs

@@ -682,7 +682,7 @@ pub fn is_dyn_method(
 /// Looks up the impl method that actually runs for the trait method `func`.
 ///
 /// Returns `func` if it's not a method defined in a trait or the lookup failed.
-pub fn lookup_impl_method(
+pub(crate) fn lookup_impl_method_query(
     db: &dyn HirDatabase,
     env: Arc<TraitEnvironment>,
     func: FunctionId,

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

@@ -32,7 +32,7 @@ pub(super) fn detect_and_exec_special_function(
         def: FunctionId,
         args: &[IntervalAndTy],
         generic_args: &Substitution,
-        locals: &Locals<'_>,
+        locals: &Locals,
         destination: Interval,
         span: MirSpan,
     ) -> Result<bool> {
@@ -168,7 +168,7 @@ fn exec_alloc_fn(
 
     fn detect_lang_function(&self, def: FunctionId) -> Option<LangItem> {
         use LangItem::*;
-        let candidate = lang_attr(self.db.upcast(), def)?;
+        let candidate = self.db.lang_attr(def.into())?;
         // We want to execute these functions with special logic
         if [PanicFmt, BeginPanic, SliceLen, DropInPlace].contains(&candidate) {
             return Some(candidate);
@@ -181,7 +181,7 @@ fn exec_lang_item(
         it: LangItem,
         generic_args: &Substitution,
         args: &[Vec<u8>],
-        locals: &Locals<'_>,
+        locals: &Locals,
         span: MirSpan,
     ) -> Result<Vec<u8>> {
         use LangItem::*;
@@ -201,7 +201,7 @@ fn exec_lang_item(
                         not_supported!("std::fmt::format not found");
                     };
                     let hir_def::resolver::ValueNs::FunctionId(format_fn) = format_fn else { not_supported!("std::fmt::format is not a function") };
-                    let message_string = self.interpret_mir(&*self.db.mir_body(format_fn.into()).map_err(|e| MirEvalError::MirLowerError(format_fn, e))?, args.cloned())?;
+                    let message_string = self.interpret_mir(self.db.mir_body(format_fn.into()).map_err(|e| MirEvalError::MirLowerError(format_fn, e))?, args.map(|x| IntervalOrOwned::Owned(x.clone())))?;
                     let addr = Address::from_bytes(&message_string[self.ptr_size()..2 * self.ptr_size()])?;
                     let size = from_bytes!(usize, message_string[2 * self.ptr_size()..]);
                     Ok(std::string::String::from_utf8_lossy(self.read_memory(addr, size)?).into_owned())
@@ -245,7 +245,7 @@ fn exec_extern_c(
         args: &[IntervalAndTy],
         _generic_args: &Substitution,
         destination: Interval,
-        locals: &Locals<'_>,
+        locals: &Locals,
         _span: MirSpan,
     ) -> Result<()> {
         match as_str {
@@ -353,7 +353,7 @@ fn exec_platform_intrinsic(
         args: &[IntervalAndTy],
         generic_args: &Substitution,
         destination: Interval,
-        locals: &Locals<'_>,
+        locals: &Locals,
         span: MirSpan,
     ) -> Result<()> {
         if let Some(name) = name.strip_prefix("simd_") {
@@ -368,7 +368,7 @@ fn exec_intrinsic(
         args: &[IntervalAndTy],
         generic_args: &Substitution,
         destination: Interval,
-        locals: &Locals<'_>,
+        locals: &Locals,
         span: MirSpan,
     ) -> Result<()> {
         if let Some(name) = name.strip_prefix("atomic_") {
@@ -873,15 +873,17 @@ fn exec_intrinsic(
                         .as_trait()
                         .and_then(|it| self.db.trait_data(it).method_by_name(&name![call_once]))
                     {
-                        return self.exec_fn_trait(
+                        self.exec_fn_trait(
                             def,
                             &args,
                             // FIXME: wrong for manual impls of `FnOnce`
                             Substitution::empty(Interner),
                             locals,
                             destination,
+                            None,
                             span,
-                        );
+                        )?;
+                        return Ok(());
                     }
                 }
                 not_supported!("FnOnce was not available for executing const_eval_select");
@@ -894,7 +896,7 @@ fn size_align_of_unsized(
         &mut self,
         ty: &Ty,
         metadata: Interval,
-        locals: &Locals<'_>,
+        locals: &Locals,
     ) -> Result<(usize, usize)> {
         Ok(match ty.kind(Interner) {
             TyKind::Str => (from_bytes!(usize, metadata.get(self)?), 1),
@@ -948,7 +950,7 @@ fn exec_atomic_intrinsic(
         args: &[IntervalAndTy],
         generic_args: &Substitution,
         destination: Interval,
-        locals: &Locals<'_>,
+        locals: &Locals,
         _span: MirSpan,
     ) -> Result<()> {
         // We are a single threaded runtime with no UB checking and no optimization, so

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

@@ -50,7 +50,7 @@ pub(super) fn exec_simd_intrinsic(
         args: &[IntervalAndTy],
         _generic_args: &Substitution,
         destination: Interval,
-        _locals: &Locals<'_>,
+        _locals: &Locals,
         _span: MirSpan,
     ) -> Result<()> {
         match name {

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

@@ -30,7 +30,7 @@ fn eval_main(db: &TestDB, file_id: FileId) -> Result<(String, String), MirEvalEr
             db.trait_environment(func_id.into()),
         )
         .map_err(|e| MirEvalError::MirLowerError(func_id.into(), e))?;
-    let (result, stdout, stderr) = interpret_mir(db, &body, false);
+    let (result, stdout, stderr) = interpret_mir(db, body, false);
     result?;
     Ok((stdout, stderr))
 }

crates/hir-ty/src/mir/lower/as_place.rs

@@ -1,7 +1,7 @@
 //! MIR lowering for places
 
 use super::*;
-use hir_def::{lang_item::lang_attr, FunctionId};
+use hir_def::FunctionId;
 use hir_expand::name;
 
 macro_rules! not_supported {
@@ -162,7 +162,7 @@ pub(super) fn lower_expr_as_place_without_adjust(
                     let is_builtin = match self.expr_ty_without_adjust(*expr).kind(Interner) {
                         TyKind::Ref(..) | TyKind::Raw(..) => true,
                         TyKind::Adt(id, _) => {
-                            if let Some(lang_item) = lang_attr(self.db.upcast(), id.0) {
+                            if let Some(lang_item) = self.db.lang_attr(id.0.into()) {
                                 lang_item == LangItem::OwnedBox
                             } else {
                                 false

crates/hir/src/lib.rs

@@ -1986,7 +1986,7 @@ pub fn eval(
                 return r;
             }
         };
-        let (result, stdout, stderr) = interpret_mir(db, &body, false);
+        let (result, stdout, stderr) = interpret_mir(db, body, false);
         let mut text = match result {
             Ok(_) => "pass".to_string(),
             Err(e) => {

crates/hir/src/source_analyzer.rs

@@ -832,7 +832,7 @@ fn resolve_impl_method_or_trait_def(
             None => return func,
         };
         let env = db.trait_environment_for_body(owner);
-        method_resolution::lookup_impl_method(db, env, func, substs).0
+        db.lookup_impl_method(env, func, substs).0
     }
 
     fn resolve_impl_const_or_trait_def(