From bc54775c9d7026b1ca616ca43fd34ff6f53d668a Mon Sep 17 00:00:00 2001
From: roife <roifewu@gmail.com>
Date: Wed, 10 Jan 2024 15:00:35 +0800
Subject: [PATCH] Preserve comments for extracted block expr in
 'extract_function' assist

---
 .../src/handlers/extract_function.rs          | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/crates/ide-assists/src/handlers/extract_function.rs b/crates/ide-assists/src/handlers/extract_function.rs
index 347a3e9ba07..3d91cb6a0f8 100644
--- a/crates/ide-assists/src/handlers/extract_function.rs
+++ b/crates/ide-assists/src/handlers/extract_function.rs
@@ -25,7 +25,7 @@ use syntax::{
         edit::{AstNodeEdit, IndentLevel},
         AstNode, HasGenericParams,
     },
-    match_ast, ted, SyntaxElement,
+    match_ast, ted, AstToken, SyntaxElement,
     SyntaxKind::{self, COMMENT},
     SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset, WalkEvent, T,
 };
@@ -1733,8 +1733,23 @@ fn make_body(
                 ast::Expr::BlockExpr(block) => {
                     // If the extracted expression is itself a block, there is no need to wrap it inside another block.
                     let block = block.dedent(old_indent);
-                    // Recreate the block for formatting consistency with other extracted functions.
-                    make::block_expr(block.statements(), block.tail_expr())
+                    let elements = block.stmt_list().map_or_else(
+                        || Either::Left(iter::empty()),
+                        |stmt_list| {
+                            let elements = stmt_list.syntax().children_with_tokens().filter_map(
+                                |node_or_token| match &node_or_token {
+                                    syntax::NodeOrToken::Node(node) => {
+                                        ast::Stmt::cast(node.clone()).map(|_| node_or_token)
+                                    }
+                                    syntax::NodeOrToken::Token(token) => {
+                                        ast::Comment::cast(token.clone()).map(|_| node_or_token)
+                                    }
+                                },
+                            );
+                            Either::Right(elements)
+                        },
+                    );
+                    make::hacky_block_expr(elements, block.tail_expr())
                 }
                 _ => {
                     let expr = expr.dedent(old_indent).indent(IndentLevel(1));