Implement prev sibling determination for `CompletionContext`

3 files changed, 114 insertions, 54 deletions

diff --git a/crates/ide_completion/src/completions/keyword.rs b/crates/ide_completion/src/completions/keyword.rs
index 662c389fe..06789b704 100644
--- a/crates/ide_completion/src/completions/keyword.rs
+++ b/crates/ide_completion/src/completions/keyword.rs
@@ -118,7 +118,7 @@ pub(crate) fn complete_expr_keyword(acc: &mut Completions, ctx: &CompletionConte
         add_keyword("let", "let ");
     }
 
-    if ctx.after_if {
+    if ctx.after_if() {
         add_keyword("else", "else {\n    $0\n}");
         add_keyword("else if", "else if $1 {\n    $0\n}");
     }
diff --git a/crates/ide_completion/src/context.rs b/crates/ide_completion/src/context.rs
index faf8469a5..8d6440cb2 100644
--- a/crates/ide_completion/src/context.rs
+++ b/crates/ide_completion/src/context.rs
@@ -17,8 +17,8 @@ use text_edit::Indel;
 
 use crate::{
     patterns::{
-        determine_location, for_is_prev2, has_prev_sibling, inside_impl_trait_block,
-        is_in_loop_body, is_match_arm, previous_token, ImmediateLocation,
+        determine_location, determine_prev_sibling, for_is_prev2, inside_impl_trait_block,
+        is_in_loop_body, is_match_arm, previous_token, ImmediateLocation, ImmediatePrevSibling,
     },
     CompletionConfig,
 };
@@ -29,12 +29,6 @@ pub(crate) enum PatternRefutability {
     Irrefutable,
 }
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
-pub(crate) enum PrevSibling {
-    Trait,
-    Impl,
-}
-
 /// `CompletionContext` is created early during completion to figure out, where
 /// exactly is the cursor, syntax-wise.
 #[derive(Debug)]
@@ -76,6 +70,7 @@ pub(crate) struct CompletionContext<'a> {
     pub(super) is_param: bool,
 
     pub(super) completion_location: Option<ImmediateLocation>,
+    pub(super) prev_sibling: Option<ImmediatePrevSibling>,
 
     /// FIXME: `ActiveParameter` is string-based, which is very very wrong
     pub(super) active_parameter: Option<ActiveParameter>,
@@ -83,7 +78,6 @@ pub(crate) struct CompletionContext<'a> {
     pub(super) is_trivial_path: bool,
     /// If not a trivial path, the prefix (qualifier).
     pub(super) path_qual: Option<ast::Path>,
-    pub(super) after_if: bool,
     /// `true` if we are a statement or a last expr in the block.
     pub(super) can_be_stmt: bool,
     /// `true` if we expect an expression at the cursor position.
@@ -107,7 +101,6 @@ pub(crate) struct CompletionContext<'a> {
 
     // keyword patterns
     pub(super) previous_token: Option<SyntaxToken>,
-    pub(super) prev_sibling: Option<PrevSibling>,
     pub(super) in_loop_body: bool,
     pub(super) is_match_arm: bool,
     pub(super) incomplete_let: bool,
@@ -173,7 +166,6 @@ impl<'a> CompletionContext<'a> {
             is_pat_or_const: None,
             is_trivial_path: false,
             path_qual: None,
-            after_if: false,
             can_be_stmt: false,
             is_expr: false,
             is_new_item: false,
@@ -308,7 +300,14 @@ impl<'a> CompletionContext<'a> {
     }
 
     pub(crate) fn has_impl_or_trait_prev_sibling(&self) -> bool {
-        self.prev_sibling.is_some()
+        matches!(
+            self.prev_sibling,
+            Some(ImmediatePrevSibling::ImplDefType) | Some(ImmediatePrevSibling::TraitDefName)
+        )
+    }
+
+    pub(crate) fn after_if(&self) -> bool {
+        matches!(self.prev_sibling, Some(ImmediatePrevSibling::IfExpr))
     }
 
     pub(crate) fn is_path_disallowed(&self) -> bool {
@@ -324,11 +323,6 @@ impl<'a> CompletionContext<'a> {
         self.previous_token = previous_token(syntax_element.clone());
         self.in_loop_body = is_in_loop_body(syntax_element.clone());
         self.is_match_arm = is_match_arm(syntax_element.clone());
-        if has_prev_sibling(syntax_element.clone(), IMPL) {
-            self.prev_sibling = Some(PrevSibling::Impl)
-        } else if has_prev_sibling(syntax_element.clone(), TRAIT) {
-            self.prev_sibling = Some(PrevSibling::Trait)
-        }
 
         self.mod_declaration_under_caret =
             find_node_at_offset::<ast::Module>(&file_with_fake_ident, offset)
@@ -468,6 +462,7 @@ impl<'a> CompletionContext<'a> {
             None => return,
         };
         self.completion_location = determine_location(&name_like);
+        self.prev_sibling = determine_prev_sibling(&name_like);
         match name_like {
             ast::NameLike::Lifetime(lifetime) => {
                 self.classify_lifetime(original_file, lifetime, offset);
@@ -656,17 +651,6 @@ impl<'a> CompletionContext<'a> {
                 })
                 .unwrap_or(false);
             self.is_expr = path.syntax().parent().and_then(ast::PathExpr::cast).is_some();
-
-            if let Some(off) = name_ref.syntax().text_range().start().checked_sub(2.into()) {
-                if let Some(if_expr) =
-                    self.sema.find_node_at_offset_with_macros::<ast::IfExpr>(original_file, off)
-                {
-                    if if_expr.syntax().text_range().end() < name_ref.syntax().text_range().start()
-                    {
-                        self.after_if = true;
-                    }
-                }
-            }
         }
 
         if let Some(field_expr) = ast::FieldExpr::cast(parent.clone()) {
diff --git a/crates/ide_completion/src/patterns.rs b/crates/ide_completion/src/patterns.rs
index f04471b57..8c4bdbed2 100644
--- a/crates/ide_completion/src/patterns.rs
+++ b/crates/ide_completion/src/patterns.rs
@@ -4,12 +4,19 @@ use syntax::{
     algo::non_trivia_sibling,
     ast::{self, LoopBodyOwner},
     match_ast, AstNode, Direction, NodeOrToken, SyntaxElement,
-    SyntaxKind::{self, *},
+    SyntaxKind::*,
     SyntaxNode, SyntaxToken, T,
 };
 
 #[cfg(test)]
 use crate::test_utils::{check_pattern_is_applicable, check_pattern_is_not_applicable};
+/// Direct parent container of the cursor position
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub(crate) enum ImmediatePrevSibling {
+    IfExpr,
+    TraitDefName,
+    ImplDefType,
+}
 
 /// Direct parent container of the cursor position
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -24,35 +31,61 @@ pub(crate) enum ImmediateLocation {
     ItemList,
 }
 
-pub(crate) fn determine_location(name_like: &ast::NameLike) -> Option<ImmediateLocation> {
-    // First walk the element we are completing up to its highest node that has the same text range
-    // as the element so that we can check in what context it immediately lies. We only do this for
-    // NameRef -> Path as that's the only thing that makes sense to being "expanded" semantically.
-    // We only wanna do this if the NameRef is the last segment of the path.
-    let node = match name_like {
-        ast::NameLike::NameRef(name_ref) => {
-            if let Some(segment) = name_ref.syntax().parent().and_then(ast::PathSegment::cast) {
-                let p = segment.parent_path();
-                if p.parent_path().is_none() {
-                    p.syntax()
-                        .ancestors()
-                        .take_while(|it| it.text_range() == p.syntax().text_range())
-                        .last()?
-                } else {
-                    return None;
+pub(crate) fn determine_prev_sibling(name_like: &ast::NameLike) -> Option<ImmediatePrevSibling> {
+    let node = maximize_name_ref(name_like)?;
+    let node = match node.parent().and_then(ast::MacroCall::cast) {
+        // When a path is being typed after the name of a trait/type of an impl it is being
+        // parsed as a macro, so when the trait/impl has a block following it an we are between the
+        // name and block the macro will attach the block to itself so maximizing fails to take
+        // that into account
+        // FIXME path expr and statement have a similar problem with attrs
+        Some(call)
+            if call.excl_token().is_none()
+                && call.token_tree().map_or(false, |t| t.l_curly_token().is_some())
+                && call.semicolon_token().is_none() =>
+        {
+            call.syntax().clone()
+        }
+        _ => node,
+    };
+    let prev_sibling = non_trivia_sibling(node.into(), Direction::Prev)?.into_node()?;
+    let res = match_ast! {
+        match prev_sibling {
+            ast::ExprStmt(it) => {
+                let node = it.expr()?.syntax().clone();
+                match_ast! {
+                    match node {
+                        ast::IfExpr(_it) => ImmediatePrevSibling::IfExpr,
+                        _ => return None,
+                    }
                 }
-            } else {
-                return None;
-            }
+            },
+            ast::Trait(it) => if it.assoc_item_list().is_none() {
+                    ImmediatePrevSibling::TraitDefName
+                } else {
+                    return None
+            },
+            ast::Impl(it) => if it.assoc_item_list().is_none()
+                && (it.for_token().is_none() || it.self_ty().is_some()) {
+                    ImmediatePrevSibling::ImplDefType
+                } else {
+                    return None
+            },
+            _ => return None,
         }
-        it @ ast::NameLike::Name(_) | it @ ast::NameLike::Lifetime(_) => it.syntax().clone(),
     };
+    Some(res)
+}
+
+pub(crate) fn determine_location(name_like: &ast::NameLike) -> Option<ImmediateLocation> {
+    let node = maximize_name_ref(name_like)?;
     let parent = match node.parent() {
         Some(parent) => match ast::MacroCall::cast(parent.clone()) {
             // When a path is being typed in an (Assoc)ItemList the parser will always emit a macro_call.
             // This is usually fine as the node expansion code above already accounts for that with
             // the ancestors call, but there is one exception to this which is that when an attribute
             // precedes it the code above will not walk the Path to the parent MacroCall as their ranges differ.
+            // FIXME path expr and statement have a similar problem
             Some(call)
                 if call.excl_token().is_none()
                     && call.token_tree().is_none()
@@ -90,6 +123,32 @@ pub(crate) fn determine_location(name_like: &ast::NameLike) -> Option<ImmediateL
     Some(res)
 }
 
+fn maximize_name_ref(name_like: &ast::NameLike) -> Option<SyntaxNode> {
+    // First walk the element we are completing up to its highest node that has the same text range
+    // as the element so that we can check in what context it immediately lies. We only do this for
+    // NameRef -> Path as that's the only thing that makes sense to being "expanded" semantically.
+    // We only wanna do this if the NameRef is the last segment of the path.
+    let node = match name_like {
+        ast::NameLike::NameRef(name_ref) => {
+            if let Some(segment) = name_ref.syntax().parent().and_then(ast::PathSegment::cast) {
+                let p = segment.parent_path();
+                if p.parent_path().is_none() {
+                    p.syntax()
+                        .ancestors()
+                        .take_while(|it| it.text_range() == p.syntax().text_range())
+                        .last()?
+                } else {
+                    return None;
+                }
+            } else {
+                return None;
+            }
+        }
+        it @ ast::NameLike::Name(_) | it @ ast::NameLike::Lifetime(_) => it.syntax().clone(),
+    };
+    Some(node)
+}
+
 #[cfg(test)]
 fn check_location(code: &str, loc: ImmediateLocation) {
     check_pattern_is_applicable(code, |e| {
@@ -192,17 +251,34 @@ fn test_for_is_prev2() {
     check_pattern_is_applicable(r"for i i$0", for_is_prev2);
 }
 
-pub(crate) fn has_prev_sibling(element: SyntaxElement, kind: SyntaxKind) -> bool {
-    previous_sibling_or_ancestor_sibling(element).filter(|it| it.kind() == kind).is_some()
+#[cfg(test)]
+fn check_prev_sibling(code: &str, sibling: impl Into<Option<ImmediatePrevSibling>>) {
+    check_pattern_is_applicable(code, |e| {
+        let name = &e.parent().and_then(ast::NameLike::cast).expect("Expected a namelike");
+        assert_eq!(determine_prev_sibling(name), sibling.into());
+        true
+    });
 }
+
 #[test]
 fn test_has_impl_as_prev_sibling() {
-    check_pattern_is_applicable(r"impl A w$0 {}", |it| has_prev_sibling(it, IMPL));
+    check_prev_sibling(r"impl A w$0 ", ImmediatePrevSibling::ImplDefType);
+    check_prev_sibling(r"impl A w$0 {}", ImmediatePrevSibling::ImplDefType);
+    check_prev_sibling(r"impl A for A w$0 ", ImmediatePrevSibling::ImplDefType);
+    check_prev_sibling(r"impl A for A w$0 {}", ImmediatePrevSibling::ImplDefType);
+    check_prev_sibling(r"impl A for w$0 {}", None);
+    check_prev_sibling(r"impl A for w$0", None);
 }
 
 #[test]
 fn test_has_trait_as_prev_sibling() {
-    check_pattern_is_applicable(r"trait A w$0 {}", |it| has_prev_sibling(it, TRAIT));
+    check_prev_sibling(r"trait A w$0 ", ImmediatePrevSibling::TraitDefName);
+    check_prev_sibling(r"trait A w$0 {}", ImmediatePrevSibling::TraitDefName);
+}
+
+#[test]
+fn test_has_if_expr_as_prev_sibling() {
+    check_prev_sibling(r"fn foo() { if true {} w$0", ImmediatePrevSibling::IfExpr);
 }
 
 pub(crate) fn is_in_loop_body(element: SyntaxElement) -> bool {