From c00f298fd26d4982e9fe092ee004facf9cef6906 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Sat, 12 Oct 2019 22:07:47 +0300
Subject: [PATCH] add syntax-tree based indents

---
 crates/ra_assists/src/assist_ctx.rs           |  1 +
 .../ra_assists/src/assists/fill_match_arms.rs | 18 ++--
 crates/ra_syntax/src/ast/edit.rs              | 91 +++++++++++++++++--
 crates/ra_syntax/src/ast/make.rs              |  6 ++
 4 files changed, 103 insertions(+), 13 deletions(-)

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 189cad7d04e..e270c5d60da 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -138,6 +138,7 @@ impl AssistBuilder {
 
     /// Replaces specified `node` of text with a given string, reindenting the
     /// string to maintain `node`'s existing indent.
+    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
     pub(crate) fn replace_node_and_indent(
         &mut self,
         node: &SyntaxNode,
diff --git a/crates/ra_assists/src/assists/fill_match_arms.rs b/crates/ra_assists/src/assists/fill_match_arms.rs
index 7335cce0940..e3f30b5defc 100644
--- a/crates/ra_assists/src/assists/fill_match_arms.rs
+++ b/crates/ra_assists/src/assists/fill_match_arms.rs
@@ -3,7 +3,7 @@
 use std::iter;
 
 use hir::{db::HirDatabase, Adt, HasSource};
-use ra_syntax::ast::{self, make, AstNode, NameOwner};
+use ra_syntax::ast::{self, edit::IndentLevel, make, AstNode, NameOwner};
 
 use crate::{Assist, AssistCtx, AssistId};
 
@@ -30,15 +30,19 @@ pub(crate) fn fill_match_arms(mut ctx: AssistCtx<impl HirDatabase>) -> Option<As
     let variant_list = enum_def.variant_list()?;
 
     ctx.add_action(AssistId("fill_match_arms"), "fill match arms", |edit| {
-        let variants = variant_list.variants();
-        let arms = variants
-            .filter_map(build_pat)
-            .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()));
-        let new_arm_list = make::match_arm_list(arms);
+        let indent_level = IndentLevel::from_node(match_arm_list.syntax());
+
+        let new_arm_list = {
+            let variants = variant_list.variants();
+            let arms = variants
+                .filter_map(build_pat)
+                .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()));
+            indent_level.increase_indent(make::match_arm_list(arms))
+        };
 
         edit.target(match_expr.syntax().text_range());
         edit.set_cursor(expr.syntax().text_range().start());
-        edit.replace_node_and_indent(match_arm_list.syntax(), new_arm_list.syntax().text());
+        edit.replace_ast(match_arm_list, new_arm_list);
     });
 
     ctx.build()
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index 03f3b5fbb00..ea92284b833 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     AstToken, Direction, InsertPosition, SmolStr, SyntaxElement,
     SyntaxKind::{ATTR, COMMENT, WHITESPACE},
-    SyntaxNode, T,
+    SyntaxNode, SyntaxToken, T,
 };
 
 impl ast::FnDef {
@@ -231,12 +231,64 @@ pub fn replace_descendants<N: AstNode, D: AstNode>(
     N::cast(new_syntax).unwrap()
 }
 
-// Note this is copy-pasted from fmt. It seems like fmt should be a separate
-// crate, but basic tree building should be this crate. However, tree building
-// might want to call into fmt...
+#[derive(Debug, Clone, Copy)]
+pub struct IndentLevel(pub u8);
+
+impl From<u8> for IndentLevel {
+    fn from(level: u8) -> IndentLevel {
+        IndentLevel(level)
+    }
+}
+
+impl IndentLevel {
+    pub fn from_node(node: &SyntaxNode) -> IndentLevel {
+        let first_token = match node.first_token() {
+            Some(it) => it,
+            None => return IndentLevel(0),
+        };
+        for ws in prev_tokens(first_token).filter_map(ast::Whitespace::cast) {
+            let text = ws.syntax().text();
+            if let Some(pos) = text.rfind('\n') {
+                let level = text[pos + 1..].chars().count() / 4;
+                return IndentLevel(level as u8);
+            }
+        }
+        IndentLevel(0)
+    }
+
+    pub fn increase_indent<N: AstNode>(self, node: N) -> N {
+        N::cast(self._increase_indent(node.syntax().clone())).unwrap()
+    }
+
+    fn _increase_indent(self, node: SyntaxNode) -> SyntaxNode {
+        let replacements: FxHashMap<SyntaxElement, SyntaxElement> = node
+            .descendants_with_tokens()
+            .filter_map(|el| el.into_token())
+            .filter_map(ast::Whitespace::cast)
+            .filter(|ws| {
+                let text = ws.syntax().text();
+                text.contains('\n')
+            })
+            .map(|ws| {
+                (
+                    ws.syntax().clone().into(),
+                    make::tokens::whitespace(&format!(
+                        "{}{:width$}",
+                        ws.syntax().text(),
+                        "",
+                        width = self.0 as usize * 4
+                    ))
+                    .into(),
+                )
+            })
+            .collect();
+        algo::replace_descendants(&node, &replacements)
+    }
+}
+
+// FIXME: replace usages with IndentLevel above
 fn leading_indent(node: &SyntaxNode) -> Option<SmolStr> {
-    let prev_tokens = std::iter::successors(node.first_token(), |token| token.prev_token());
-    for token in prev_tokens {
+    for token in prev_tokens(node.first_token()?) {
         if let Some(ws) = ast::Whitespace::cast(token.clone()) {
             let ws_text = ws.text();
             if let Some(pos) = ws_text.rfind('\n') {
@@ -250,6 +302,10 @@ fn leading_indent(node: &SyntaxNode) -> Option<SmolStr> {
     None
 }
 
+fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
+    iter::successors(Some(token), |token| token.prev_token())
+}
+
 #[must_use]
 fn insert_children<N: AstNode>(
     parent: &N,
@@ -269,3 +325,26 @@ fn replace_children<N: AstNode>(
     let new_syntax = algo::replace_children(parent.syntax(), to_replace, &mut to_insert);
     N::cast(new_syntax).unwrap()
 }
+
+#[test]
+fn test_increase_indent() {
+    let arm_list = {
+        let arm = make::match_arm(iter::once(make::placeholder_pat().into()), make::expr_unit());
+        make::match_arm_list(vec![arm.clone(), arm].into_iter())
+    };
+    assert_eq!(
+        arm_list.syntax().to_string(),
+        "{
+    _ => (),
+    _ => (),
+}"
+    );
+    let indented = IndentLevel(2).increase_indent(arm_list);
+    assert_eq!(
+        indented.syntax().to_string(),
+        "{
+            _ => (),
+            _ => (),
+        }"
+    );
+}
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 287a40bee76..143835172af 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -160,6 +160,12 @@ pub mod tokens {
             .unwrap()
     }
 
+    pub fn whitespace(text: &str) -> SyntaxToken {
+        assert!(text.trim().is_empty());
+        let sf = SourceFile::parse(text).ok().unwrap();
+        sf.syntax().first_child_or_token().unwrap().into_token().unwrap()
+    }
+
     pub fn single_newline() -> SyntaxToken {
         SOURCE_FILE
             .tree()