fix: make make_body respect comments in extract_function

1 files changed, 63 insertions, 4 deletions

diff --git a/crates/ide-assists/src/handlers/extract_function.rs b/crates/ide-assists/src/handlers/extract_function.rs
index 0483cfdc646..4ee9d4638bb 100644
--- a/crates/ide-assists/src/handlers/extract_function.rs
+++ b/crates/ide-assists/src/handlers/extract_function.rs
@@ -1799,7 +1799,6 @@ fn make_body(
                 })
                 .collect::<Vec<SyntaxElement>>();
             let tail_expr = tail_expr.map(|expr| expr.dedent(old_indent).indent(body_indent));
-
             make::hacky_block_expr_with_comments(elements, tail_expr)
         }
     };
@@ -1860,9 +1859,29 @@ fn with_default_tail_expr(block: ast::BlockExpr, tail_expr: ast::Expr) -> ast::B
 }
 
 fn with_tail_expr(block: ast::BlockExpr, tail_expr: ast::Expr) -> ast::BlockExpr {
-    let stmt_tail = block.tail_expr().map(|expr| make::expr_stmt(expr).into());
-    let stmts = block.statements().chain(stmt_tail);
-    make::block_expr(stmts, Some(tail_expr))
+    let stmt_tail_opt: Option<ast::Stmt> =
+        block.tail_expr().map(|expr| make::expr_stmt(expr).into());
+
+    let mut elements: Vec<SyntaxElement> = vec![];
+
+    block.statements().for_each(|stmt| {
+        elements.push(syntax::NodeOrToken::Node(stmt.syntax().clone()));
+    });
+
+    if let Some(stmt_list) = block.stmt_list() {
+        stmt_list.syntax().children_with_tokens().for_each(|node_or_token| {
+            match &node_or_token {
+                syntax::NodeOrToken::Token(_) => elements.push(node_or_token),
+                _ => (),
+            };
+        });
+    }
+
+    if let Some(stmt_tail) = stmt_tail_opt {
+        elements.push(syntax::NodeOrToken::Node(stmt_tail.syntax().clone()));
+    }
+
+    make::hacky_block_expr_with_comments(elements, Some(tail_expr))
 }
 
 fn format_type(ty: &hir::Type, ctx: &AssistContext<'_>, module: hir::Module) -> String {
@@ -5747,4 +5766,44 @@ fn $0fun_name() -> Option<()> {
 "#,
         );
     }
+
+    #[test]
+    fn non_tail_expr_with_comment_of_tail_expr_loop() {
+        check_assist(
+            extract_function,
+            r#"
+pub fn f() {
+    loop {
+        $0// A comment
+        if true {
+            continue;
+        }$0
+        if false {
+            break;
+        }
+    }
+}
+"#,
+            r#"
+pub fn f() {
+    loop {
+        if let ControlFlow::Break(_) = fun_name() {
+            continue;
+        }
+        if false {
+            break;
+        }
+    }
+}
+
+fn $0fun_name() -> ControlFlow<()> {
+    // A comment
+    if true {
+        return ControlFlow::Break(());
+    }
+    ControlFlow::Continue(())
+}
+"#,
+        );
+    }
 }