Indent block expressions on enter

crates/ide/src/typing/on_enter.rs

@@ -1,13 +1,20 @@
 //! Handles the `Enter` key press. At the momently, this only continues
 //! comments, but should handle indent some time in the future as well.
 
-use ide_db::base_db::{FilePosition, SourceDatabase};
+use std::sync::Arc;
+
 use ide_db::RootDatabase;
+use ide_db::{
+    base_db::{FilePosition, SourceDatabase},
+    line_index::LineIndex,
+    LineIndexDatabase,
+};
 use syntax::{
-    ast::{self, AstToken},
+    algo::find_node_at_offset,
+    ast::{self, edit::IndentLevel, AstToken},
     AstNode, SmolStr, SourceFile,
     SyntaxKind::*,
-    SyntaxToken, TextRange, TextSize, TokenAtOffset,
+    SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset,
 };
 
 use text_edit::TextEdit;
@@ -19,6 +26,7 @@ use text_edit::TextEdit;
 // - kbd:[Enter] inside triple-slash comments automatically inserts `///`
 // - kbd:[Enter] in the middle or after a trailing space in `//` inserts `//`
 // - kbd:[Enter] inside `//!` doc comments automatically inserts `//!`
+// - kbd:[Enter] after `{` indents contents and closing `}` of single-line block
 //
 // This action needs to be assigned to shortcut explicitly.
 //
@@ -38,25 +46,42 @@ use text_edit::TextEdit;
 pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<TextEdit> {
     let parse = db.parse(position.file_id);
     let file = parse.tree();
-    let comment = file
+    let token = file.syntax().token_at_offset(position.offset).left_biased()?;
-        .syntax()
-        .token_at_offset(position.offset)
-        .left_biased()
-        .and_then(ast::Comment::cast)?;
 
+    if let Some(comment) = ast::Comment::cast(token.clone()) {
+        return on_enter_in_comment(&comment, &file, position.offset);
+    }
+
+    if token.kind() == L_CURLY {
+        // Typing enter after the `{` of a block expression, where the `}` is on the same line
+        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
+            .and_then(|block| on_enter_in_block(db, block, position))
+        {
+            return Some(edit);
+        }
+    }
+
+    None
+}
+
+fn on_enter_in_comment(
+    comment: &ast::Comment,
+    file: &ast::SourceFile,
+    offset: TextSize,
+) -> Option<TextEdit> {
     if comment.kind().shape.is_block() {
         return None;
     }
 
     let prefix = comment.prefix();
     let comment_range = comment.syntax().text_range();
-    if position.offset < comment_range.start() + TextSize::of(prefix) {
+    if offset < comment_range.start() + TextSize::of(prefix) {
         return None;
     }
 
     let mut remove_trailing_whitespace = false;
     // Continuing single-line non-doc comments (like this one :) ) is annoying
-    if prefix == "//" && comment_range.end() == position.offset {
+    if prefix == "//" && comment_range.end() == offset {
         if comment.text().ends_with(' ') {
             cov_mark::hit!(continues_end_of_line_comment_with_space);
             remove_trailing_whitespace = true;
@@ -70,14 +95,50 @@ pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<Text
     let delete = if remove_trailing_whitespace {
         let trimmed_len = comment.text().trim_end().len() as u32;
         let trailing_whitespace_len = comment.text().len() as u32 - trimmed_len;
-        TextRange::new(position.offset - TextSize::from(trailing_whitespace_len), position.offset)
+        TextRange::new(offset - TextSize::from(trailing_whitespace_len), offset)
     } else {
-        TextRange::empty(position.offset)
+        TextRange::empty(offset)
     };
     let edit = TextEdit::replace(delete, inserted);
     Some(edit)
 }
 
+fn on_enter_in_block(
+    db: &RootDatabase,
+    block: ast::BlockExpr,
+    position: FilePosition,
+) -> Option<TextEdit> {
+    let contents = block_contents(&block)?;
+
+    let line_index: Arc<LineIndex> = db.line_index(position.file_id);
+    let (open, close) = (block.l_curly_token()?, block.r_curly_token()?);
+    let start = line_index.line_col(open.text_range().start()).line;
+    let end = line_index.line_col(close.text_range().end()).line;
+    if start != end {
+        return None;
+    }
+
+    let indent = IndentLevel::from_node(block.syntax());
+    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
+    edit.union(TextEdit::insert(contents.text_range().end(), format!("\n{}", indent))).ok()?;
+    Some(edit)
+}
+
+fn block_contents(block: &ast::BlockExpr) -> Option<SyntaxNode> {
+    let mut node = block.tail_expr().map(|e| e.syntax().clone());
+
+    for stmt in block.statements() {
+        if node.is_some() {
+            // More than 1 node in the block
+            return None;
+        }
+
+        node = Some(stmt.syntax().clone());
+    }
+
+    node
+}
+
 fn followed_by_comment(comment: &ast::Comment) -> bool {
     let ws = match comment.syntax().next_token().and_then(ast::Whitespace::cast) {
         Some(it) => it,
@@ -296,4 +357,143 @@ fn main() {
 ",
         );
     }
+
+    #[test]
+    fn indents_fn_body_block() {
+        do_check(
+            r#"
+fn f() {$0()}
+        "#,
+            r#"
+fn f() {
+    $0()
+}
+        "#,
+        );
+    }
+
+    #[test]
+    fn indents_block_expr() {
+        do_check(
+            r#"
+fn f() {
+    let x = {$0()};
+}
+        "#,
+            r#"
+fn f() {
+    let x = {
+        $0()
+    };
+}
+        "#,
+        );
+    }
+
+    #[test]
+    fn indents_match_arm() {
+        do_check(
+            r#"
+fn f() {
+    match 6 {
+        1 => {$0f()},
+        _ => (),
+    }
+}
+        "#,
+            r#"
+fn f() {
+    match 6 {
+        1 => {
+            $0f()
+        },
+        _ => (),
+    }
+}
+        "#,
+        );
+    }
+
+    #[test]
+    fn indents_block_with_statement() {
+        do_check(
+            r#"
+fn f() {$0a = b}
+        "#,
+            r#"
+fn f() {
+    $0a = b
+}
+        "#,
+        );
+        do_check(
+            r#"
+fn f() {$0fn f() {}}
+        "#,
+            r#"
+fn f() {
+    $0fn f() {}
+}
+        "#,
+        );
+    }
+
+    #[test]
+    fn indents_nested_blocks() {
+        do_check(
+            r#"
+fn f() {$0{}}
+        "#,
+            r#"
+fn f() {
+    $0{}
+}
+        "#,
+        );
+    }
+
+    #[test]
+    fn does_not_indent_empty_block() {
+        do_check_noop(
+            r#"
+fn f() {$0}
+        "#,
+        );
+        do_check_noop(
+            r#"
+fn f() {{$0}}
+        "#,
+        );
+    }
+
+    #[test]
+    fn does_not_indent_block_with_too_much_content() {
+        do_check_noop(
+            r#"
+fn f() {$0 a = b; ()}
+        "#,
+        );
+        do_check_noop(
+            r#"
+fn f() {$0 a = b; a = b; }
+        "#,
+        );
+    }
+
+    #[test]
+    fn does_not_indent_multiline_block() {
+        do_check_noop(
+            r#"
+fn f() {$0
+}
+        "#,
+        );
+        do_check_noop(
+            r#"
+fn f() {$0
+
+}
+        "#,
+        );
+    }
 }