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

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

Author: ChayimFriedman2

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

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,
     ) {

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
+"#]],
+        );
+    }
 }

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,