From ae73628f6bf0ee3ac6f571b447972f4f37a39103 Mon Sep 17 00:00:00 2001
From: bvanjoi <bohan-zhang@foxmail.com>
Date: Wed, 4 Jan 2023 22:10:17 +0800
Subject: [PATCH] fix: keep whitespace in extract function handler

---
 .../src/handlers/extract_function.rs            |  9 +++++----
 crates/syntax/src/ast/make.rs                   | 17 ++++++++++++-----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/crates/ide-assists/src/handlers/extract_function.rs b/crates/ide-assists/src/handlers/extract_function.rs
index 4ee9d4638bbd5..94614360c7292 100644
--- a/crates/ide-assists/src/handlers/extract_function.rs
+++ b/crates/ide-assists/src/handlers/extract_function.rs
@@ -1799,7 +1799,8 @@ 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)
+
+            make::hacky_block_expr(elements, tail_expr)
         }
     };
 
@@ -1881,7 +1882,7 @@ fn with_tail_expr(block: ast::BlockExpr, tail_expr: ast::Expr) -> ast::BlockExpr
         elements.push(syntax::NodeOrToken::Node(stmt_tail.syntax().clone()));
     }
 
-    make::hacky_block_expr_with_comments(elements, Some(tail_expr))
+    make::hacky_block_expr(elements, Some(tail_expr))
 }
 
 fn format_type(ty: &hir::Type, ctx: &AssistContext<'_>, module: hir::Module) -> String {
@@ -4978,9 +4979,8 @@ fn $0fun_name() {
         );
     }
 
-    // FIXME: we do want to preserve whitespace
     #[test]
-    fn extract_function_does_not_preserve_whitespace() {
+    fn extract_function_does_preserve_whitespace() {
         check_assist(
             extract_function,
             r#"
@@ -4999,6 +4999,7 @@ fn func() {
 
 fn $0fun_name() {
     let a = 0;
+
     let x = 0;
 }
 "#,
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index 8c26009add2bb..f17b7d8557df3 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -339,10 +339,10 @@ pub fn tail_only_block_expr(tail_expr: ast::Expr) -> ast::BlockExpr {
 }
 
 /// Ideally this function wouldn't exist since it involves manual indenting.
-/// It differs from `make::block_expr` by also supporting comments.
+/// It differs from `make::block_expr` by also supporting comments and whitespace.
 ///
 /// FIXME: replace usages of this with the mutable syntax tree API
-pub fn hacky_block_expr_with_comments(
+pub fn hacky_block_expr(
     elements: impl IntoIterator<Item = crate::SyntaxElement>,
     tail_expr: Option<ast::Expr>,
 ) -> ast::BlockExpr {
@@ -350,10 +350,17 @@ pub fn hacky_block_expr_with_comments(
     for node_or_token in elements.into_iter() {
         match node_or_token {
             rowan::NodeOrToken::Node(n) => format_to!(buf, "    {n}\n"),
-            rowan::NodeOrToken::Token(t) if t.kind() == SyntaxKind::COMMENT => {
-                format_to!(buf, "    {t}\n")
+            rowan::NodeOrToken::Token(t) => {
+                let kind = t.kind();
+                if kind == SyntaxKind::COMMENT {
+                    format_to!(buf, "    {t}\n")
+                } else if kind == SyntaxKind::WHITESPACE {
+                    let content = t.text().trim_matches(|c| c != '\n');
+                    if content.len() >= 1 {
+                        format_to!(buf, "{}", &content[1..])
+                    }
+                }
             }
-            _ => (),
         }
     }
     if let Some(tail_expr) = tail_expr {