Add smart completions that skip `await` or `iter()` and `into_iter()`

E.g. complete `await.foo()`.

7 files changed, 230 insertions, 72 deletions

diff --git a/src/tools/rust-analyzer/crates/hir/src/lib.rs b/src/tools/rust-analyzer/crates/hir/src/lib.rs
index 29649cac94c..0d2728bb613 100644
--- a/src/tools/rust-analyzer/crates/hir/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/lib.rs
@@ -58,8 +58,7 @@ use hir_def::{
     CrateRootModuleId, DefWithBodyId, EnumId, EnumVariantId, ExternCrateId, FunctionId,
     GenericDefId, GenericParamId, HasModule, ImplId, InTypeConstId, ItemContainerId,
     LifetimeParamId, LocalFieldId, Lookup, MacroExpander, MacroId, ModuleId, StaticId, StructId,
-    SyntheticSyntax, TraitAliasId, TraitId, TupleId, TypeAliasId, TypeOrConstParamId, TypeParamId,
-    UnionId,
+    SyntheticSyntax, TraitAliasId, TupleId, TypeAliasId, TypeOrConstParamId, TypeParamId, UnionId,
 };
 use hir_expand::{
     attrs::collect_attrs, proc_macro::ProcMacroKind, AstId, MacroCallKind, RenderedExpandError,
@@ -128,7 +127,7 @@ pub use {
         ImportPathConfig,
         // FIXME: This is here since some queries take it as input that are used
         // outside of hir.
-        ModuleDefId,
+        {ModuleDefId, TraitId},
     },
     hir_expand::{
         attrs::{Attr, AttrId},
@@ -4749,6 +4748,14 @@ impl Type {
         Some((self.derived(ty.clone()), m))
     }
 
+    pub fn add_reference(&self, mutability: Mutability) -> Type {
+        let ty_mutability = match mutability {
+            Mutability::Shared => hir_ty::Mutability::Not,
+            Mutability::Mut => hir_ty::Mutability::Mut,
+        };
+        self.derived(TyKind::Ref(ty_mutability, error_lifetime(), self.ty.clone()).intern(Interner))
+    }
+
     pub fn is_slice(&self) -> bool {
         matches!(self.ty.kind(Interner), TyKind::Slice(..))
     }
@@ -4804,9 +4811,9 @@ impl Type {
     }
 
     /// Checks that particular type `ty` implements `std::future::IntoFuture` or
-    /// `std::future::Future`.
+    /// `std::future::Future` and returns the `Output` associated type.
     /// This function is used in `.await` syntax completion.
-    pub fn impls_into_future(&self, db: &dyn HirDatabase) -> bool {
+    pub fn into_future_output(&self, db: &dyn HirDatabase) -> Option<Type> {
         let trait_ = db
             .lang_item(self.env.krate, LangItem::IntoFutureIntoFuture)
             .and_then(|it| {
@@ -4818,16 +4825,18 @@ impl Type {
             .or_else(|| {
                 let future_trait = db.lang_item(self.env.krate, LangItem::Future)?;
                 future_trait.as_trait()
-            });
-
-        let trait_ = match trait_ {
-            Some(it) => it,
-            None => return false,
-        };
+            })?;
 
         let canonical_ty =
             Canonical { value: self.ty.clone(), binders: CanonicalVarKinds::empty(Interner) };
-        method_resolution::implements_trait(&canonical_ty, db, &self.env, trait_)
+        if !method_resolution::implements_trait_unique(&canonical_ty, db, &self.env, trait_) {
+            return None;
+        }
+
+        let output_assoc_type = db
+            .trait_data(trait_)
+            .associated_type_by_name(&Name::new_symbol_root(sym::Output.clone()))?;
+        self.normalize_trait_assoc_type(db, &[], output_assoc_type.into())
     }
 
     /// This does **not** resolve `IntoFuture`, only `Future`.
@@ -4846,6 +4855,26 @@ impl Type {
         self.normalize_trait_assoc_type(db, &[], iterator_item.into())
     }
 
+    pub fn into_iterator_iter(self, db: &dyn HirDatabase) -> Option<Type> {
+        let trait_ = db.lang_item(self.env.krate, LangItem::IntoIterIntoIter).and_then(|it| {
+            let into_iter_fn = it.as_function()?;
+            let assoc_item = as_assoc_item(db, AssocItem::Function, into_iter_fn)?;
+            let into_iter_trait = assoc_item.container_or_implemented_trait(db)?;
+            Some(into_iter_trait.id)
+        })?;
+
+        let canonical_ty =
+            Canonical { value: self.ty.clone(), binders: CanonicalVarKinds::empty(Interner) };
+        if !method_resolution::implements_trait_unique(&canonical_ty, db, &self.env, trait_) {
+            return None;
+        }
+
+        let into_iter_assoc_type = db
+            .trait_data(trait_)
+            .associated_type_by_name(&Name::new_symbol_root(sym::IntoIter.clone()))?;
+        self.normalize_trait_assoc_type(db, &[], into_iter_assoc_type.into())
+    }
+
     /// Checks that particular type `ty` implements `std::ops::FnOnce`.
     ///
     /// This function can be used to check if a particular type is callable, since FnOnce is a
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/completions.rs b/src/tools/rust-analyzer/crates/ide-completion/src/completions.rs
index 414627fbaba..40669c65c57 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/completions.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/completions.rs
@@ -329,7 +329,7 @@ impl Completions {
         ctx: &CompletionContext<'_>,
         dot_access: &DotAccess,
         func: hir::Function,
-        receiver: Option<hir::Name>,
+        receiver: Option<SmolStr>,
         local_name: Option<hir::Name>,
     ) {
         if !ctx.check_stability(Some(&func.attrs(ctx.db))) {
@@ -475,7 +475,7 @@ impl Completions {
         &mut self,
         ctx: &CompletionContext<'_>,
         dot_access: &DotAccess,
-        receiver: Option<hir::Name>,
+        receiver: Option<SmolStr>,
         field: hir::Field,
         ty: &hir::Type,
     ) {
@@ -533,7 +533,7 @@ impl Completions {
     pub(crate) fn add_tuple_field(
         &mut self,
         ctx: &CompletionContext<'_>,
-        receiver: Option<hir::Name>,
+        receiver: Option<SmolStr>,
         field: usize,
         ty: &hir::Type,
     ) {
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/completions/dot.rs b/src/tools/rust-analyzer/crates/ide-completion/src/completions/dot.rs
index 8d44dbab606..3a991e007af 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/completions/dot.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/completions/dot.rs
@@ -2,7 +2,7 @@
 
 use std::ops::ControlFlow;
 
-use hir::{sym, HasContainer, ItemContainer, MethodCandidateCallback, Name};
+use hir::{HasContainer, ItemContainer, MethodCandidateCallback, Name};
 use ide_db::FxHashSet;
 use syntax::SmolStr;
 
@@ -25,8 +25,13 @@ pub(crate) fn complete_dot(
         _ => return,
     };
 
+    let is_field_access = matches!(dot_access.kind, DotAccessKind::Field { .. });
+    let is_method_access_with_parens =
+        matches!(dot_access.kind, DotAccessKind::Method { has_parens: true });
+    let traits_in_scope = ctx.traits_in_scope();
+
     // Suggest .await syntax for types that implement Future trait
-    if receiver_ty.impls_into_future(ctx.db) {
+    if let Some(future_output) = receiver_ty.into_future_output(ctx.db) {
         let mut item = CompletionItem::new(
             CompletionItemKind::Keyword,
             ctx.source_range(),
@@ -35,11 +40,37 @@ pub(crate) fn complete_dot(
         );
         item.detail("expr.await");
         item.add_to(acc, ctx.db);
-    }
 
-    let is_field_access = matches!(dot_access.kind, DotAccessKind::Field { .. });
-    let is_method_access_with_parens =
-        matches!(dot_access.kind, DotAccessKind::Method { has_parens: true });
+        // Completions that skip `.await`, e.g. `.await.foo()`.
+        let dot_access_kind = match &dot_access.kind {
+            DotAccessKind::Field { receiver_is_ambiguous_float_literal: _ } => {
+                DotAccessKind::Field { receiver_is_ambiguous_float_literal: false }
+            }
+            it @ DotAccessKind::Method { .. } => *it,
+        };
+        let dot_access = DotAccess {
+            receiver: dot_access.receiver.clone(),
+            receiver_ty: Some(hir::TypeInfo { original: future_output.clone(), adjusted: None }),
+            kind: dot_access_kind,
+            ctx: dot_access.ctx,
+        };
+        complete_fields(
+            acc,
+            ctx,
+            &future_output,
+            |acc, field, ty| {
+                acc.add_field(ctx, &dot_access, Some(SmolStr::new_static("await")), field, &ty)
+            },
+            |acc, field, ty| {
+                acc.add_tuple_field(ctx, Some(SmolStr::new_static("await")), field, &ty)
+            },
+            is_field_access,
+            is_method_access_with_parens,
+        );
+        complete_methods(ctx, &future_output, &traits_in_scope, |func| {
+            acc.add_method(ctx, &dot_access, func, Some(SmolStr::new_static("await")), None)
+        });
+    }
 
     complete_fields(
         acc,
@@ -50,8 +81,41 @@ pub(crate) fn complete_dot(
         is_field_access,
         is_method_access_with_parens,
     );
+    complete_methods(ctx, receiver_ty, &traits_in_scope, |func| {
+        acc.add_method(ctx, dot_access, func, None, None)
+    });
 
-    complete_methods(ctx, receiver_ty, |func| acc.add_method(ctx, dot_access, func, None, None));
+    // Checking for the existence of `iter()` is complicated in our setup, because we need to substitute
+    // its return type, so we instead check for `<&Self as IntoIterator>::IntoIter`.
+    let iter = receiver_ty
+        .strip_references()
+        .add_reference(hir::Mutability::Shared)
+        .into_iterator_iter(ctx.db)
+        .map(|ty| (ty, SmolStr::new_static("iter()")))
+        .or_else(|| {
+            receiver_ty
+                .clone()
+                .into_iterator_iter(ctx.db)
+                .map(|ty| (ty, SmolStr::new_static("into_iter()")))
+        });
+    if let Some((iter, iter_sym)) = iter {
+        // Skip iterators, e.g. complete `.iter().filter_map()`.
+        let dot_access_kind = match &dot_access.kind {
+            DotAccessKind::Field { receiver_is_ambiguous_float_literal: _ } => {
+                DotAccessKind::Field { receiver_is_ambiguous_float_literal: false }
+            }
+            it @ DotAccessKind::Method { .. } => *it,
+        };
+        let dot_access = DotAccess {
+            receiver: dot_access.receiver.clone(),
+            receiver_ty: Some(hir::TypeInfo { original: iter.clone(), adjusted: None }),
+            kind: dot_access_kind,
+            ctx: dot_access.ctx,
+        };
+        complete_methods(ctx, &iter, &traits_in_scope, |func| {
+            acc.add_method(ctx, &dot_access, func, Some(iter_sym.clone()), None)
+        });
+    }
 }
 
 pub(crate) fn complete_undotted_self(
@@ -94,18 +158,16 @@ pub(crate) fn complete_undotted_self(
                         in_breakable: expr_ctx.in_breakable,
                     },
                 },
-                Some(Name::new_symbol_root(sym::self_.clone())),
+                Some(SmolStr::new_static("self")),
                 field,
                 &ty,
             )
         },
-        |acc, field, ty| {
-            acc.add_tuple_field(ctx, Some(Name::new_symbol_root(sym::self_.clone())), field, &ty)
-        },
+        |acc, field, ty| acc.add_tuple_field(ctx, Some(SmolStr::new_static("self")), field, &ty),
         true,
         false,
     );
-    complete_methods(ctx, &ty, |func| {
+    complete_methods(ctx, &ty, &ctx.traits_in_scope(), |func| {
         acc.add_method(
             ctx,
             &DotAccess {
@@ -118,7 +180,7 @@ pub(crate) fn complete_undotted_self(
                 },
             },
             func,
-            Some(Name::new_symbol_root(sym::self_.clone())),
+            Some(SmolStr::new_static("self")),
             None,
         )
     });
@@ -160,6 +222,7 @@ fn complete_fields(
 fn complete_methods(
     ctx: &CompletionContext<'_>,
     receiver: &hir::Type,
+    traits_in_scope: &FxHashSet<hir::TraitId>,
     f: impl FnMut(hir::Function),
 ) {
     struct Callback<'a, F> {
@@ -205,7 +268,7 @@ fn complete_methods(
     receiver.iterate_method_candidates_split_inherent(
         ctx.db,
         &ctx.scope,
-        &ctx.traits_in_scope(),
+        traits_in_scope,
         Some(ctx.module),
         None,
         Callback { ctx, f, seen_methods: FxHashSet::default() },
@@ -1306,4 +1369,73 @@ fn baz() {
             "#]],
         );
     }
+
+    #[test]
+    fn skip_iter() {
+        check_no_kw(
+            r#"
+        //- minicore: iterator
+        fn foo() {
+            [].$0
+        }
+        "#,
+            expect![[r#"
+                me clone() (as Clone)                                       fn(&self) -> Self
+                me into_iter() (as IntoIterator) fn(self) -> <Self as IntoIterator>::IntoIter
+            "#]],
+        );
+        check_no_kw(
+            r#"
+//- minicore: iterator
+struct MyIntoIter;
+impl IntoIterator for MyIntoIter {
+    type Item = ();
+    type IntoIter = MyIterator;
+    fn into_iter(self) -> Self::IntoIter {
+        MyIterator
+    }
+}
+
+struct MyIterator;
+impl Iterator for MyIterator {
+    type Item = ();
+    fn next(&mut self) -> Self::Item {}
+}
+
+fn foo() {
+    MyIntoIter.$0
+}
+"#,
+            expect![[r#"
+                me into_iter() (as IntoIterator)                fn(self) -> <Self as IntoIterator>::IntoIter
+                me into_iter().by_ref() (as Iterator)                             fn(&mut self) -> &mut Self
+                me into_iter().into_iter() (as IntoIterator)    fn(self) -> <Self as IntoIterator>::IntoIter
+                me into_iter().next() (as Iterator)        fn(&mut self) -> Option<<Self as Iterator>::Item>
+                me into_iter().nth(…) (as Iterator) fn(&mut self, usize) -> Option<<Self as Iterator>::Item>
+            "#]],
+        );
+    }
+
+    #[test]
+    fn skip_await() {
+        check_no_kw(
+            r#"
+//- minicore: future
+struct Foo;
+impl Foo {
+    fn foo(self) {}
+}
+
+async fn foo() -> Foo { Foo }
+
+async fn bar() {
+    foo().$0
+}
+"#,
+            expect![[r#"
+    me await.foo()                                                                      fn(self)
+    me into_future() (use core::future::IntoFuture) fn(self) -> <Self as IntoFuture>::IntoFuture
+"#]],
+        );
+    }
 }
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/context.rs b/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
index 3705e2c73d6..461b1dc036c 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
@@ -390,7 +390,7 @@ pub(crate) struct DotAccess {
     pub(crate) ctx: DotAccessExprCtx,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub(crate) enum DotAccessKind {
     Field {
         /// True if the receiver is an integer and there is no ident in the original file after it yet
@@ -402,7 +402,7 @@ pub(crate) enum DotAccessKind {
     },
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct DotAccessExprCtx {
     pub(crate) in_block_expr: bool,
     pub(crate) in_breakable: BreakableKind,
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/render.rs b/src/tools/rust-analyzer/crates/ide-completion/src/render.rs
index eb5d136b8b3..0990a32e38c 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/render.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/render.rs
@@ -18,7 +18,7 @@ use ide_db::{
     imports::import_assets::LocatedImport,
     RootDatabase, SnippetCap, SymbolKind,
 };
-use syntax::{ast, format_smolstr, AstNode, Edition, SmolStr, SyntaxKind, TextRange, ToSmolStr};
+use syntax::{ast, format_smolstr, AstNode, SmolStr, SyntaxKind, TextRange, ToSmolStr};
 
 use crate::{
     context::{DotAccess, DotAccessKind, PathCompletionCtx, PathKind, PatternContext},
@@ -122,7 +122,7 @@ impl<'a> RenderContext<'a> {
 pub(crate) fn render_field(
     ctx: RenderContext<'_>,
     dot_access: &DotAccess,
-    receiver: Option<hir::Name>,
+    receiver: Option<SmolStr>,
     field: hir::Field,
     ty: &hir::Type,
 ) -> CompletionItem {
@@ -136,7 +136,7 @@ pub(crate) fn render_field(
     let mut item = CompletionItem::new(
         SymbolKind::Field,
         ctx.source_range(),
-        field_with_receiver(db, receiver.as_ref(), &name, ctx.completion.edition),
+        field_with_receiver(receiver.as_deref(), &name),
         ctx.completion.edition,
     );
     item.set_relevance(CompletionRelevance {
@@ -158,8 +158,7 @@ pub(crate) fn render_field(
 
         builder.replace(
             ctx.source_range(),
-            field_with_receiver(db, receiver.as_ref(), &escaped_name, ctx.completion.edition)
-                .into(),
+            field_with_receiver(receiver.as_deref(), &escaped_name).into(),
         );
 
         let expected_fn_type =
@@ -183,12 +182,7 @@ pub(crate) fn render_field(
 
         item.text_edit(builder.finish());
     } else {
-        item.insert_text(field_with_receiver(
-            db,
-            receiver.as_ref(),
-            &escaped_name,
-            ctx.completion.edition,
-        ));
+        item.insert_text(field_with_receiver(receiver.as_deref(), &escaped_name));
     }
     if let Some(receiver) = &dot_access.receiver {
         if let Some(original) = ctx.completion.sema.original_ast_node(receiver.clone()) {
@@ -201,33 +195,21 @@ pub(crate) fn render_field(
     item.build(db)
 }
 
-fn field_with_receiver(
-    db: &RootDatabase,
-    receiver: Option<&hir::Name>,
-    field_name: &str,
-    edition: Edition,
-) -> SmolStr {
-    receiver.map_or_else(
-        || field_name.into(),
-        |receiver| format_smolstr!("{}.{field_name}", receiver.display(db, edition)),
-    )
+fn field_with_receiver(receiver: Option<&str>, field_name: &str) -> SmolStr {
+    receiver
+        .map_or_else(|| field_name.into(), |receiver| format_smolstr!("{}.{field_name}", receiver))
 }
 
 pub(crate) fn render_tuple_field(
     ctx: RenderContext<'_>,
-    receiver: Option<hir::Name>,
+    receiver: Option<SmolStr>,
     field: usize,
     ty: &hir::Type,
 ) -> CompletionItem {
     let mut item = CompletionItem::new(
         SymbolKind::Field,
         ctx.source_range(),
-        field_with_receiver(
-            ctx.db(),
-            receiver.as_ref(),
-            &field.to_string(),
-            ctx.completion.edition,
-        ),
+        field_with_receiver(receiver.as_deref(), &field.to_string()),
         ctx.completion.edition,
     );
     item.detail(ty.display(ctx.db(), ctx.completion.edition).to_string())
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs b/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
index a859d79e243..3b579f5c0bf 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
@@ -23,7 +23,7 @@ use crate::{
 #[derive(Debug)]
 enum FuncKind<'ctx> {
     Function(&'ctx PathCompletionCtx),
-    Method(&'ctx DotAccess, Option<hir::Name>),
+    Method(&'ctx DotAccess, Option<SmolStr>),
 }
 
 pub(crate) fn render_fn(
@@ -39,7 +39,7 @@ pub(crate) fn render_fn(
 pub(crate) fn render_method(
     ctx: RenderContext<'_>,
     dot_access: &DotAccess,
-    receiver: Option<hir::Name>,
+    receiver: Option<SmolStr>,
     local_name: Option<hir::Name>,
     func: hir::Function,
 ) -> Builder {
@@ -59,16 +59,8 @@ fn render(
 
     let (call, escaped_call) = match &func_kind {
         FuncKind::Method(_, Some(receiver)) => (
-            format_smolstr!(
-                "{}.{}",
-                receiver.unescaped().display(ctx.db()),
-                name.unescaped().display(ctx.db())
-            ),
-            format_smolstr!(
-                "{}.{}",
-                receiver.display(ctx.db(), completion.edition),
-                name.display(ctx.db(), completion.edition)
-            ),
+            format_smolstr!("{}.{}", receiver, name.unescaped().display(ctx.db())),
+            format_smolstr!("{}.{}", receiver, name.display(ctx.db(), completion.edition)),
         ),
         _ => (
             name.unescaped().display(db).to_smolstr(),
diff --git a/src/tools/rust-analyzer/crates/test-utils/src/minicore.rs b/src/tools/rust-analyzer/crates/test-utils/src/minicore.rs
index 4a2346193b4..4462c2ce1e1 100644
--- a/src/tools/rust-analyzer/crates/test-utils/src/minicore.rs
+++ b/src/tools/rust-analyzer/crates/test-utils/src/minicore.rs
@@ -1510,7 +1510,7 @@ pub mod iter {
             impl<T, const N: usize> IntoIterator for [T; N] {
                 type Item = T;
                 type IntoIter = IntoIter<T, N>;
-                fn into_iter(self) -> I {
+                fn into_iter(self) -> Self::IntoIter {
                     IntoIter { data: self, range: IndexRange { start: 0, end: loop {} } }
                 }
             }
@@ -1520,6 +1520,29 @@ pub mod iter {
                     loop {}
                 }
             }
+            pub struct Iter<'a, T> {
+                slice: &'a [T],
+            }
+            impl<'a, T> IntoIterator for &'a [T; N] {
+                type Item = &'a T;
+                type IntoIter = Iter<'a, T>;
+                fn into_iter(self) -> Self::IntoIter {
+                    loop {}
+                }
+            }
+            impl<'a, T> IntoIterator for &'a [T] {
+                type Item = &'a T;
+                type IntoIter = Iter<'a, T>;
+                fn into_iter(self) -> Self::IntoIter {
+                    loop {}
+                }
+            }
+            impl<'a, T> Iterator for Iter<'a, T> {
+                type Item = &'a T;
+                fn next(&mut self) -> Option<T> {
+                    loop {}
+                }
+            }
         }
         pub use self::collect::IntoIterator;
     }