1 files changed, 150 insertions, 9 deletions

diff --git a/src/tools/rust-analyzer/crates/hir/src/source_analyzer.rs b/src/tools/rust-analyzer/crates/hir/src/source_analyzer.rs
index 1eb51b20c..342912b67 100644
--- a/src/tools/rust-analyzer/crates/hir/src/source_analyzer.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/source_analyzer.rs
@@ -25,15 +25,21 @@ use hir_def::{
     Lookup, ModuleDefId, VariantId,
 };
 use hir_expand::{
-    builtin_fn_macro::BuiltinFnLikeExpander, hygiene::Hygiene, name::AsName, HirFileId, InFile,
+    builtin_fn_macro::BuiltinFnLikeExpander,
+    hygiene::Hygiene,
+    mod_path::path,
+    name,
+    name::{AsName, Name},
+    HirFileId, InFile,
 };
 use hir_ty::{
     diagnostics::{
         record_literal_missing_fields, record_pattern_missing_fields, unsafe_expressions,
         UnsafeExpr,
     },
-    method_resolution, Adjust, Adjustment, AutoBorrow, InferenceResult, Interner, Substitution,
-    TyExt, TyKind, TyLoweringContext,
+    method_resolution::{self, lang_names_for_bin_op},
+    Adjust, Adjustment, AutoBorrow, InferenceResult, Interner, Substitution, Ty, TyExt, TyKind,
+    TyLoweringContext,
 };
 use itertools::Itertools;
 use smallvec::SmallVec;
@@ -134,11 +140,19 @@ impl SourceAnalyzer {
     ) -> Option<InFile<ast::Expr>> {
         let macro_file = self.body_source_map()?.node_macro_file(expr.as_ref())?;
         let expanded = db.parse_or_expand(macro_file)?;
-
-        let res = match ast::MacroCall::cast(expanded.clone()) {
-            Some(call) => self.expand_expr(db, InFile::new(macro_file, call))?,
-            _ => InFile::new(macro_file, ast::Expr::cast(expanded)?),
+        let res = if let Some(stmts) = ast::MacroStmts::cast(expanded.clone()) {
+            match stmts.expr()? {
+                ast::Expr::MacroExpr(mac) => {
+                    self.expand_expr(db, InFile::new(macro_file, mac.macro_call()?))?
+                }
+                expr => InFile::new(macro_file, expr),
+            }
+        } else if let Some(call) = ast::MacroCall::cast(expanded.clone()) {
+            self.expand_expr(db, InFile::new(macro_file, call))?
+        } else {
+            InFile::new(macro_file, ast::Expr::cast(expanded)?)
         };
+
         Some(res)
     }
 
@@ -255,8 +269,111 @@ impl SourceAnalyzer {
     ) -> Option<FunctionId> {
         let expr_id = self.expr_id(db, &call.clone().into())?;
         let (f_in_trait, substs) = self.infer.as_ref()?.method_resolution(expr_id)?;
-        let f_in_impl = self.resolve_impl_method(db, f_in_trait, &substs);
-        f_in_impl.or(Some(f_in_trait))
+
+        Some(self.resolve_impl_method_or_trait_def(db, f_in_trait, &substs))
+    }
+
+    pub(crate) fn resolve_await_to_poll(
+        &self,
+        db: &dyn HirDatabase,
+        await_expr: &ast::AwaitExpr,
+    ) -> Option<FunctionId> {
+        let mut ty = self.ty_of_expr(db, &await_expr.expr()?.into())?.clone();
+
+        let into_future_trait = self
+            .resolver
+            .resolve_known_trait(db.upcast(), &path![core::future::IntoFuture])
+            .map(Trait::from);
+
+        if let Some(into_future_trait) = into_future_trait {
+            let type_ = Type::new_with_resolver(db, &self.resolver, ty.clone());
+            if type_.impls_trait(db, into_future_trait, &[]) {
+                let items = into_future_trait.items(db);
+                let into_future_type = items.into_iter().find_map(|item| match item {
+                    AssocItem::TypeAlias(alias)
+                        if alias.name(db) == hir_expand::name![IntoFuture] =>
+                    {
+                        Some(alias)
+                    }
+                    _ => None,
+                })?;
+                let future_trait = type_.normalize_trait_assoc_type(db, &[], into_future_type)?;
+                ty = future_trait.ty;
+            }
+        }
+
+        let poll_fn = db
+            .lang_item(self.resolver.krate(), hir_expand::name![poll].to_smol_str())?
+            .as_function()?;
+        let substs = hir_ty::TyBuilder::subst_for_def(db, poll_fn).push(ty.clone()).build();
+        Some(self.resolve_impl_method_or_trait_def(db, poll_fn, &substs))
+    }
+
+    pub(crate) fn resolve_prefix_expr(
+        &self,
+        db: &dyn HirDatabase,
+        prefix_expr: &ast::PrefixExpr,
+    ) -> Option<FunctionId> {
+        let lang_item_name = match prefix_expr.op_kind()? {
+            ast::UnaryOp::Deref => name![deref],
+            ast::UnaryOp::Not => name![not],
+            ast::UnaryOp::Neg => name![neg],
+        };
+        let ty = self.ty_of_expr(db, &prefix_expr.expr()?.into())?;
+
+        let op_fn = self.lang_trait_fn(db, &lang_item_name, &lang_item_name)?;
+        let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
+
+        Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
+    }
+
+    pub(crate) fn resolve_index_expr(
+        &self,
+        db: &dyn HirDatabase,
+        index_expr: &ast::IndexExpr,
+    ) -> Option<FunctionId> {
+        let base_ty = self.ty_of_expr(db, &index_expr.base()?.into())?;
+        let index_ty = self.ty_of_expr(db, &index_expr.index()?.into())?;
+
+        let lang_item_name = name![index];
+
+        let op_fn = self.lang_trait_fn(db, &lang_item_name, &lang_item_name)?;
+        let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn)
+            .push(base_ty.clone())
+            .push(index_ty.clone())
+            .build();
+        Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
+    }
+
+    pub(crate) fn resolve_bin_expr(
+        &self,
+        db: &dyn HirDatabase,
+        binop_expr: &ast::BinExpr,
+    ) -> Option<FunctionId> {
+        let op = binop_expr.op_kind()?;
+        let lhs = self.ty_of_expr(db, &binop_expr.lhs()?.into())?;
+        let rhs = self.ty_of_expr(db, &binop_expr.rhs()?.into())?;
+
+        let op_fn = lang_names_for_bin_op(op)
+            .and_then(|(name, lang_item)| self.lang_trait_fn(db, &lang_item, &name))?;
+        let substs =
+            hir_ty::TyBuilder::subst_for_def(db, op_fn).push(lhs.clone()).push(rhs.clone()).build();
+
+        Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
+    }
+
+    pub(crate) fn resolve_try_expr(
+        &self,
+        db: &dyn HirDatabase,
+        try_expr: &ast::TryExpr,
+    ) -> Option<FunctionId> {
+        let ty = self.ty_of_expr(db, &try_expr.expr()?.into())?;
+
+        let op_fn =
+            db.lang_item(self.resolver.krate(), name![branch].to_smol_str())?.as_function()?;
+        let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
+
+        Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
     }
 
     pub(crate) fn resolve_field(
@@ -281,6 +398,7 @@ impl SourceAnalyzer {
         let local = if field.name_ref().is_some() {
             None
         } else {
+            // Shorthand syntax, resolve to the local
             let path = ModPath::from_segments(PathKind::Plain, once(local_name.clone()));
             match self.resolver.resolve_path_in_value_ns_fully(db.upcast(), &path) {
                 Some(ValueNs::LocalBinding(pat_id)) => {
@@ -666,6 +784,29 @@ impl SourceAnalyzer {
         let fun_data = db.function_data(func);
         method_resolution::lookup_impl_method(self_ty, db, trait_env, impled_trait, &fun_data.name)
     }
+
+    fn resolve_impl_method_or_trait_def(
+        &self,
+        db: &dyn HirDatabase,
+        func: FunctionId,
+        substs: &Substitution,
+    ) -> FunctionId {
+        self.resolve_impl_method(db, func, substs).unwrap_or(func)
+    }
+
+    fn lang_trait_fn(
+        &self,
+        db: &dyn HirDatabase,
+        lang_trait: &Name,
+        method_name: &Name,
+    ) -> Option<FunctionId> {
+        db.trait_data(db.lang_item(self.resolver.krate(), lang_trait.to_smol_str())?.as_trait()?)
+            .method_by_name(method_name)
+    }
+
+    fn ty_of_expr(&self, db: &dyn HirDatabase, expr: &ast::Expr) -> Option<&Ty> {
+        self.infer.as_ref()?.type_of_expr.get(self.expr_id(db, &expr)?)
+    }
 }
 
 fn scope_for(