Merge pull request #19129 from ChayimFriedman2/snippet-macro

fix: Fix postfix completions inside macros

Author: Veykril

crates/ide-completion/src/completions/postfix.rs

@@ -2,17 +2,18 @@
 
 mod format_like;
 
-use hir::ItemInNs;
-use ide_db::text_edit::TextEdit;
+use base_db::SourceDatabase;
+use hir::{ItemInNs, Semantics};
 use ide_db::{
     documentation::{Documentation, HasDocs},
     imports::insert_use::ImportScope,
+    text_edit::TextEdit,
     ty_filter::TryEnum,
-    SnippetCap,
+    RootDatabase, SnippetCap,
 };
 use stdx::never;
 use syntax::{
-    ast::{self, make, AstNode, AstToken},
+    ast::{self, AstNode, AstToken},
     SyntaxKind::{BLOCK_EXPR, EXPR_STMT, FOR_EXPR, IF_EXPR, LOOP_EXPR, STMT_LIST, WHILE_EXPR},
     TextRange, TextSize,
 };
@@ -48,7 +49,8 @@ pub(crate) fn complete_postfix(
     };
     let expr_ctx = &dot_access.ctx;
 
-    let receiver_text = get_receiver_text(dot_receiver, receiver_is_ambiguous_float_literal);
+    let receiver_text =
+        get_receiver_text(&ctx.sema, dot_receiver, receiver_is_ambiguous_float_literal);
 
     let cap = match ctx.config.snippet_cap {
         Some(it) => it,
@@ -172,13 +174,15 @@ pub(crate) fn complete_postfix(
     // The rest of the postfix completions create an expression that moves an argument,
     // so it's better to consider references now to avoid breaking the compilation
 
-    let (dot_receiver, node_to_replace_with) = include_references(dot_receiver);
-    let receiver_text =
-        get_receiver_text(&node_to_replace_with, receiver_is_ambiguous_float_literal);
-    let postfix_snippet = match build_postfix_snippet_builder(ctx, cap, &dot_receiver) {
-        Some(it) => it,
-        None => return,
-    };
+    let (dot_receiver_including_refs, prefix) = include_references(dot_receiver);
+    let mut receiver_text =
+        get_receiver_text(&ctx.sema, dot_receiver, receiver_is_ambiguous_float_literal);
+    receiver_text.insert_str(0, &prefix);
+    let postfix_snippet =
+        match build_postfix_snippet_builder(ctx, cap, &dot_receiver_including_refs) {
+            Some(it) => it,
+            None => return,
+        };
 
     if !ctx.config.snippets.is_empty() {
         add_custom_postfix_completions(acc, ctx, &postfix_snippet, &receiver_text);
@@ -222,7 +226,7 @@ pub(crate) fn complete_postfix(
     postfix_snippet("call", "function(expr)", &format!("${{1}}({receiver_text})"))
         .add_to(acc, ctx.db);
 
-    if let Some(parent) = dot_receiver.syntax().parent().and_then(|p| p.parent()) {
+    if let Some(parent) = dot_receiver_including_refs.syntax().parent().and_then(|p| p.parent()) {
         if matches!(parent.kind(), STMT_LIST | EXPR_STMT) {
             postfix_snippet("let", "let", &format!("let $0 = {receiver_text};"))
                 .add_to(acc, ctx.db);
@@ -231,9 +235,9 @@ pub(crate) fn complete_postfix(
         }
     }
 
-    if let ast::Expr::Literal(literal) = dot_receiver.clone() {
+    if let ast::Expr::Literal(literal) = dot_receiver_including_refs.clone() {
         if let Some(literal_text) = ast::String::cast(literal.token()) {
-            add_format_like_completions(acc, ctx, &dot_receiver, cap, &literal_text);
+            add_format_like_completions(acc, ctx, &dot_receiver_including_refs, cap, &literal_text);
         }
     }
 
@@ -260,14 +264,20 @@ pub(crate) fn complete_postfix(
     }
 }
 
-fn get_receiver_text(receiver: &ast::Expr, receiver_is_ambiguous_float_literal: bool) -> String {
-    let mut text = if receiver_is_ambiguous_float_literal {
-        let text = receiver.syntax().text();
-        let without_dot = ..text.len() - TextSize::of('.');
-        text.slice(without_dot).to_string()
-    } else {
-        receiver.to_string()
+fn get_receiver_text(
+    sema: &Semantics<'_, RootDatabase>,
+    receiver: &ast::Expr,
+    receiver_is_ambiguous_float_literal: bool,
+) -> String {
+    // Do not just call `receiver.to_string()`, as that will mess up whitespaces inside macros.
+    let Some(mut range) = sema.original_range_opt(receiver.syntax()) else {
+        return receiver.to_string();
     };
+    if receiver_is_ambiguous_float_literal {
+        range.range = TextRange::at(range.range.start(), range.range.len() - TextSize::of('.'))
+    }
+    let file_text = sema.db.file_text(range.file_id.file_id());
+    let mut text = file_text[range.range].to_owned();
 
     // The receiver texts should be interpreted as-is, as they are expected to be
     // normal Rust expressions.
@@ -284,15 +294,15 @@ fn escape_snippet_bits(text: &mut String) {
     stdx::replace(text, '$', "\\$");
 }
 
-fn include_references(initial_element: &ast::Expr) -> (ast::Expr, ast::Expr) {
+fn include_references(initial_element: &ast::Expr) -> (ast::Expr, String) {
     let mut resulting_element = initial_element.clone();
 
     while let Some(field_expr) = resulting_element.syntax().parent().and_then(ast::FieldExpr::cast)
     {
         resulting_element = ast::Expr::from(field_expr);
     }
 
-    let mut new_element_opt = initial_element.clone();
+    let mut prefix = String::new();
 
     while let Some(parent_deref_element) =
         resulting_element.syntax().parent().and_then(ast::PrefixExpr::cast)
@@ -303,7 +313,7 @@ fn include_references(initial_element: &ast::Expr) -> (ast::Expr, ast::Expr) {
 
         resulting_element = ast::Expr::from(parent_deref_element);
 
-        new_element_opt = make::expr_prefix(syntax::T![*], new_element_opt).into();
+        prefix.insert(0, '*');
     }
 
     if let Some(first_ref_expr) = resulting_element.syntax().parent().and_then(ast::RefExpr::cast) {
@@ -317,15 +327,15 @@ fn include_references(initial_element: &ast::Expr) -> (ast::Expr, ast::Expr) {
             let exclusive = parent_ref_element.mut_token().is_some();
             resulting_element = ast::Expr::from(parent_ref_element);
 
-            new_element_opt = make::expr_ref(new_element_opt, exclusive);
+            prefix.insert_str(0, if exclusive { "&mut " } else { "&" });
         }
     } else {
         // If we do not find any ref expressions, restore
         // all the progress of tree climbing
         resulting_element = initial_element.clone();
     }
 
-    (resulting_element, new_element_opt)
+    (resulting_element, prefix)
 }
 
 fn build_postfix_snippet_builder<'ctx>(
@@ -901,4 +911,31 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn inside_macro() {
+        check_edit(
+            "box",
+            r#"
+macro_rules! assert {
+    ( $it:expr $(,)? ) => { $it };
+}
+
+fn foo() {
+    let a = true;
+    assert!(if a == false { true } else { false }.$0);
+}
+        "#,
+            r#"
+macro_rules! assert {
+    ( $it:expr $(,)? ) => { $it };
+}
+
+fn foo() {
+    let a = true;
+    assert!(Box::new(if a == false { true } else { false }));
+}
+        "#,
+        );
+    }
 }