1 files changed, 219 insertions, 11 deletions

diff --git a/crates/ide/src/typing/on_enter.rs b/crates/ide/src/typing/on_enter.rs
index 9144681bf..7d2db201a 100644
--- a/crates/ide/src/typing/on_enter.rs
+++ b/crates/ide/src/typing/on_enter.rs
@@ -4,10 +4,11 @@
 use ide_db::base_db::{FilePosition, SourceDatabase};
 use ide_db::RootDatabase;
 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;
@@ -18,6 +19,8 @@ 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.
 //
@@ -37,25 +40,43 @@ 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
-        .syntax()
-        .token_at_offset(position.offset)
-        .left_biased()
-        .and_then(ast::Comment::cast)?;
+    let token = file.syntax().token_at_offset(position.offset).left_biased()?;
 
+    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(block, position))
+        {
+            cov_mark::hit!(indent_block_contents);
+            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;
@@ -69,14 +90,42 @@ 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(block: ast::BlockExpr, position: FilePosition) -> Option<TextEdit> {
+    let contents = block_contents(&block)?;
+
+    if block.syntax().text().contains_char('\n') {
+        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,
@@ -187,6 +236,25 @@ fn foo() {
     }
 
     #[test]
+    fn continues_another_doc_comment() {
+        do_check(
+            r#"
+fn main() {
+    //! Documentation for$0 on enter
+    let x = 1 + 1;
+}
+"#,
+            r#"
+fn main() {
+    //! Documentation for
+    //! $0 on enter
+    let x = 1 + 1;
+}
+"#,
+        );
+    }
+
+    #[test]
     fn continues_code_comment_in_the_middle_of_line() {
         do_check(
             r"
@@ -276,4 +344,144 @@ fn main() {
 ",
         );
     }
+
+    #[test]
+    fn indents_fn_body_block() {
+        cov_mark::check!(indent_block_contents);
+        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
+
+}
+        "#,
+        );
+    }
 }