Precompute expected type during completion

Author: matklad

crates/ra_ide/src/completion/complete_unqualified_path.rs

@@ -4,7 +4,7 @@ use hir::ScopeDef;
 use test_utils::tested_by;
 
 use crate::completion::{CompletionContext, Completions};
-use hir::{Adt, ModuleDef};
+use hir::{Adt, ModuleDef, Type};
 use ra_syntax::AstNode;
 
 pub(super) fn complete_unqualified_path(acc: &mut Completions, ctx: &CompletionContext) {
@@ -15,7 +15,9 @@ pub(super) fn complete_unqualified_path(acc: &mut Completions, ctx: &CompletionC
         return;
     }
 
-    complete_enum_variants(acc, ctx);
+    if let Some(ty) = &ctx.expected_type {
+        complete_enum_variants(acc, ctx, ty);
+    }
 
     if ctx.is_pat_binding_or_const {
         return;
@@ -34,26 +36,24 @@ pub(super) fn complete_unqualified_path(acc: &mut Completions, ctx: &CompletionC
     });
 }
 
-fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext) {
-    if let Some(ty) = ctx.expected_type_of(&ctx.token.parent()) {
-        if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
-            let variants = enum_data.variants(ctx.db);
-
-            let module = if let Some(module) = ctx.scope().module() {
-                // Compute path from the completion site if available.
-                module
-            } else {
-                // Otherwise fall back to the enum's definition site.
-                enum_data.module(ctx.db)
-            };
-
-            for variant in variants {
-                if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
-                    // Variants with trivial paths are already added by the existing completion logic,
-                    // so we should avoid adding these twice
-                    if path.segments.len() > 1 {
-                        acc.add_enum_variant(ctx, variant, Some(path.to_string()));
-                    }
+fn complete_enum_variants(acc: &mut Completions, ctx: &CompletionContext, ty: &Type) {
+    if let Some(Adt::Enum(enum_data)) = ty.as_adt() {
+        let variants = enum_data.variants(ctx.db);
+
+        let module = if let Some(module) = ctx.scope().module() {
+            // Compute path from the completion site if available.
+            module
+        } else {
+            // Otherwise fall back to the enum's definition site.
+            enum_data.module(ctx.db)
+        };
+
+        for variant in variants {
+            if let Some(path) = module.find_use_path(ctx.db, ModuleDef::from(variant)) {
+                // Variants with trivial paths are already added by the existing completion logic,
+                // so we should avoid adding these twice
+                if path.segments.len() > 1 {
+                    acc.add_enum_variant(ctx, variant, Some(path.to_string()));
                 }
             }
         }

crates/ra_ide/src/completion/completion_context.rs

@@ -5,7 +5,7 @@ use ra_db::SourceDatabase;
 use ra_ide_db::RootDatabase;
 use ra_syntax::{
     algo::{find_covering_element, find_node_at_offset},
-    ast, AstNode,
+    ast, match_ast, AstNode,
     SyntaxKind::*,
     SyntaxNode, SyntaxToken, TextRange, TextSize,
 };
@@ -26,6 +26,7 @@ pub(crate) struct CompletionContext<'a> {
     /// The token before the cursor, in the macro-expanded file.
     pub(super) token: SyntaxToken,
     pub(super) krate: Option<hir::Crate>,
+    pub(super) expected_type: Option<Type>,
     pub(super) name_ref_syntax: Option<ast::NameRef>,
     pub(super) function_syntax: Option<ast::FnDef>,
     pub(super) use_item_syntax: Option<ast::UseItem>,
@@ -93,6 +94,7 @@ impl<'a> CompletionContext<'a> {
             token,
             offset: position.offset,
             krate,
+            expected_type: None,
             name_ref_syntax: None,
             function_syntax: None,
             use_item_syntax: None,
@@ -175,23 +177,30 @@ impl<'a> CompletionContext<'a> {
         self.sema.scope_at_offset(&self.token.parent(), self.offset)
     }
 
-    pub(crate) fn expected_type_of(&self, node: &SyntaxNode) -> Option<Type> {
-        for ancestor in node.ancestors() {
-            if let Some(pat) = ast::Pat::cast(ancestor.clone()) {
-                return self.sema.type_of_pat(&pat);
-            } else if let Some(expr) = ast::Expr::cast(ancestor) {
-                return self.sema.type_of_expr(&expr);
-            }
-        }
-        None
-    }
-
     fn fill(
         &mut self,
         original_file: &SyntaxNode,
         file_with_fake_ident: SyntaxNode,
         offset: TextSize,
     ) {
+        // FIXME: this is wrong in at least two cases:
+        //  * when there's no token `foo(<|>)`
+        //  * when there is a token, but it happens to have type of it's own
+        self.expected_type = self
+            .token
+            .ancestors()
+            .find_map(|node| {
+                let ty = match_ast! {
+                    match node {
+                        ast::Pat(it) => self.sema.type_of_pat(&it),
+                        ast::Expr(it) => self.sema.type_of_expr(&it),
+                        _ => return None,
+                    }
+                };
+                Some(ty)
+            })
+            .flatten();
+
         // First, let's try to complete a reference to some declaration.
         if let Some(name_ref) = find_node_at_offset::<ast::NameRef>(&file_with_fake_ident, offset) {
             // Special case, `trait T { fn foo(i_am_a_name_ref) {} }`.

crates/ra_ide/src/completion/presentation.rs

@@ -351,7 +351,7 @@ impl Builder {
         }
 
         // Don't add parentheses if the expected type is some function reference.
-        if let Some(ty) = ctx.expected_type_of(&ctx.token.parent()) {
+        if let Some(ty) = &ctx.expected_type {
             if ty.is_fn() {
                 return self;
             }