Move `CompletionContext::function_def` into `PathKind::Expr`

4 files changed, 82 insertions, 39 deletions

diff --git a/crates/ide-completion/src/completions/dot.rs b/crates/ide-completion/src/completions/dot.rs
index a315d616d55..da26c2ad682 100644
--- a/crates/ide-completion/src/completions/dot.rs
+++ b/crates/ide-completion/src/completions/dot.rs
@@ -46,27 +46,26 @@ pub(crate) fn complete_undotted_self(
     if !ctx.config.enable_self_on_the_fly {
         return;
     }
-    match path_ctx {
-        PathCompletionCtx { qualified: Qualified::No, kind: PathKind::Expr { .. }, .. }
-            if path_ctx.is_trivial_path() && ctx.qualifier_ctx.none() => {}
+    let self_param = match path_ctx {
+        PathCompletionCtx {
+            qualified: Qualified::No,
+            kind: PathKind::Expr { self_param: Some(self_param), .. },
+            ..
+        } if path_ctx.is_trivial_path() && ctx.qualifier_ctx.none() => self_param,
         _ => return,
-    }
+    };
 
-    if let Some(func) = ctx.function_def.as_ref().and_then(|fn_| ctx.sema.to_def(fn_)) {
-        if let Some(self_) = func.self_param(ctx.db) {
-            let ty = self_.ty(ctx.db);
-            complete_fields(
-                acc,
-                ctx,
-                &ty,
-                |acc, field, ty| acc.add_field(ctx, Some(hir::known::SELF_PARAM), field, &ty),
-                |acc, field, ty| acc.add_tuple_field(ctx, Some(hir::known::SELF_PARAM), field, &ty),
-            );
-            complete_methods(ctx, &ty, |func| {
-                acc.add_method(ctx, func, Some(hir::known::SELF_PARAM), None)
-            });
-        }
-    }
+    let ty = self_param.ty(ctx.db);
+    complete_fields(
+        acc,
+        ctx,
+        &ty,
+        |acc, field, ty| acc.add_field(ctx, Some(hir::known::SELF_PARAM), field, &ty),
+        |acc, field, ty| acc.add_tuple_field(ctx, Some(hir::known::SELF_PARAM), field, &ty),
+    );
+    complete_methods(ctx, &ty, |func| {
+        acc.add_method(ctx, func, Some(hir::known::SELF_PARAM), None)
+    });
 }
 
 fn complete_fields(
diff --git a/crates/ide-completion/src/completions/expr.rs b/crates/ide-completion/src/completions/expr.rs
index 6152ccb711f..ecc1442bfc7 100644
--- a/crates/ide-completion/src/completions/expr.rs
+++ b/crates/ide-completion/src/completions/expr.rs
@@ -14,7 +14,9 @@ pub(crate) fn complete_expr_path(
     path_ctx: &PathCompletionCtx,
 ) {
     let _p = profile::span("complete_expr_path");
-
+    if !ctx.qualifier_ctx.none() {
+        return;
+    }
     let (
         qualified,
         in_block_expr,
@@ -23,6 +25,7 @@ pub(crate) fn complete_expr_path(
         after_if_expr,
         wants_mut_token,
         in_condition,
+        ty,
     ) = match path_ctx {
         &PathCompletionCtx {
             kind:
@@ -33,10 +36,12 @@ pub(crate) fn complete_expr_path(
                     in_condition,
                     ref ref_expr_parent,
                     ref is_func_update,
+                    ref innermost_ret_ty,
+                    ..
                 },
             ref qualified,
             ..
-        } if ctx.qualifier_ctx.none() => (
+        } => (
             qualified,
             in_block_expr,
             in_loop_body,
@@ -44,6 +49,7 @@ pub(crate) fn complete_expr_path(
             after_if_expr,
             ref_expr_parent.as_ref().map(|it| it.mut_token().is_none()).unwrap_or(false),
             in_condition,
+            innermost_ret_ty,
         ),
         _ => return,
     };
@@ -252,10 +258,10 @@ pub(crate) fn complete_expr_path(
                     }
                 }
 
-                if let Some(fn_def) = &ctx.function_def {
+                if let Some(ty) = ty {
                     add_keyword(
                         "return",
-                        match (in_block_expr, fn_def.ret_type().is_some()) {
+                        match (in_block_expr, ty.is_unit()) {
                             (true, true) => "return ;",
                             (true, false) => "return;",
                             (false, true) => "return $0",
diff --git a/crates/ide-completion/src/context.rs b/crates/ide-completion/src/context.rs
index d73cb6034d2..7df9b4921af 100644
--- a/crates/ide-completion/src/context.rs
+++ b/crates/ide-completion/src/context.rs
@@ -95,6 +95,8 @@ pub(super) enum PathKind {
         in_condition: bool,
         ref_expr_parent: Option<ast::RefExpr>,
         is_func_update: Option<ast::RecordExpr>,
+        self_param: Option<hir::SelfParam>,
+        innermost_ret_ty: Option<hir::Type>,
     },
     Type {
         location: TypeLocation,
@@ -317,9 +319,6 @@ pub(crate) struct CompletionContext<'a> {
     /// The expected type of what we are completing.
     pub(super) expected_type: Option<Type>,
 
-    /// The parent function of the cursor position if it exists.
-    // FIXME: This probably doesn't belong here
-    pub(super) function_def: Option<ast::Fn>,
     /// The parent impl of the cursor position if it exists.
     // FIXME: This probably doesn't belong here
     pub(super) impl_def: Option<ast::Impl>,
@@ -500,7 +499,6 @@ impl<'a> CompletionContext<'a> {
             module,
             expected_name: None,
             expected_type: None,
-            function_def: None,
             impl_def: None,
             incomplete_let: false,
             previous_token: None,
diff --git a/crates/ide-completion/src/context/analysis.rs b/crates/ide-completion/src/context/analysis.rs
index fc6ee70a0bc..4f65ae402e2 100644
--- a/crates/ide-completion/src/context/analysis.rs
+++ b/crates/ide-completion/src/context/analysis.rs
@@ -404,18 +404,6 @@ impl<'a> CompletionContext<'a> {
                 ast::Item::Impl(impl_) => Some(impl_),
                 _ => None,
             });
-        self.function_def = self
-            .sema
-            .token_ancestors_with_macros(self.token.clone())
-            .take_while(|it| {
-                it.kind() != SyntaxKind::SOURCE_FILE && it.kind() != SyntaxKind::MODULE
-            })
-            .filter_map(ast::Item::cast)
-            .take(2)
-            .find_map(|it| match it {
-                ast::Item::Fn(fn_) => Some(fn_),
-                _ => None,
-            });
 
         match name_like {
             ast::NameLike::Lifetime(lifetime) => {
@@ -727,6 +715,56 @@ impl<'a> CompletionContext<'a> {
             let after_if_expr = after_if_expr(it.clone());
             let ref_expr_parent =
                 path.as_single_name_ref().and_then(|_| it.parent()).and_then(ast::RefExpr::cast);
+            let (innermost_ret_ty, self_param) = {
+                let find_ret_ty = |it: SyntaxNode| {
+                    if let Some(item) = ast::Item::cast(it.clone()) {
+                        match item {
+                            ast::Item::Fn(f) => {
+                                Some(sema.to_def(&f).map(|it| it.ret_type(sema.db)))
+                            }
+                            ast::Item::MacroCall(_) => None,
+                            _ => Some(None),
+                        }
+                    } else {
+                        let expr = ast::Expr::cast(it)?;
+                        let callable = match expr {
+                            // FIXME
+                            // ast::Expr::BlockExpr(b) if b.async_token().is_some() || b.try_token().is_some() => sema.type_of_expr(b),
+                            ast::Expr::ClosureExpr(_) => sema.type_of_expr(&expr),
+                            _ => return None,
+                        };
+                        Some(
+                            callable
+                                .and_then(|c| c.adjusted().as_callable(sema.db))
+                                .map(|it| it.return_type()),
+                        )
+                    }
+                };
+                let find_fn_self_param = |it| match it {
+                    ast::Item::Fn(fn_) => {
+                        Some(sema.to_def(&fn_).and_then(|it| it.self_param(sema.db)))
+                    }
+                    ast::Item::MacroCall(_) => None,
+                    _ => Some(None),
+                };
+
+                match dbg!(find_node_in_file_compensated(original_file, &expr)) {
+                    Some(it) => {
+                        let innermost_ret_ty = sema
+                            .ancestors_with_macros(it.syntax().clone())
+                            .find_map(find_ret_ty)
+                            .flatten();
+
+                        let self_param = sema
+                            .ancestors_with_macros(it.syntax().clone())
+                            .filter_map(ast::Item::cast)
+                            .find_map(find_fn_self_param)
+                            .flatten();
+                        (innermost_ret_ty, self_param)
+                    }
+                    None => (None, None),
+                }
+            };
             let is_func_update = func_update_record(it);
             let in_condition = is_in_condition(&expr);
 
@@ -737,6 +775,8 @@ impl<'a> CompletionContext<'a> {
                 in_condition,
                 ref_expr_parent,
                 is_func_update,
+                innermost_ret_ty,
+                self_param,
             }
         };
         let make_path_kind_type = |ty: ast::Type| {