Fix postfix completions inside macros

Previously the receiver text was taken directly from the AST, which in macros is missing trivia, leading to corruption (or just unintended replacement of user code).

Now we upmap the range, and extract the original file text in it.

Author: ChayimFriedman2

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