7 files changed, 339 insertions, 36 deletions

diff --git a/crates/ide/src/diagnostics.rs b/crates/ide/src/diagnostics.rs
index dd42116a7..0ace80a1e 100644
--- a/crates/ide/src/diagnostics.rs
+++ b/crates/ide/src/diagnostics.rs
@@ -20,7 +20,7 @@ use itertools::Itertools;
 use rustc_hash::FxHashSet;
 use syntax::{
     ast::{self, AstNode},
-    SyntaxNode, SyntaxNodePtr, TextRange,
+    SyntaxNode, SyntaxNodePtr, TextRange, TextSize,
 };
 use text_edit::TextEdit;
 use unlinked_file::UnlinkedFile;
@@ -159,14 +159,16 @@ pub(crate) fn diagnostics(
             );
         })
         .on::<UnlinkedFile, _>(|d| {
+            // Limit diagnostic to the first few characters in the file. This matches how VS Code
+            // renders it with the full span, but on other editors, and is less invasive.
+            let range = sema.diagnostics_display_range(d.display_source()).range;
+            let range = range.intersect(TextRange::up_to(TextSize::of("..."))).unwrap_or(range);
+
             // Override severity and mark as unused.
             res.borrow_mut().push(
-                Diagnostic::hint(
-                    sema.diagnostics_display_range(d.display_source()).range,
-                    d.message(),
-                )
-                .with_fix(d.fix(&sema))
-                .with_code(Some(d.code())),
+                Diagnostic::hint(range, d.message())
+                    .with_fix(d.fix(&sema))
+                    .with_code(Some(d.code())),
             );
         })
         .on::<hir::diagnostics::UnresolvedProcMacro, _>(|d| {
diff --git a/crates/ide/src/expand_macro.rs b/crates/ide/src/expand_macro.rs
index 9eeabbeda..be0ee03bf 100644
--- a/crates/ide/src/expand_macro.rs
+++ b/crates/ide/src/expand_macro.rs
@@ -1,3 +1,5 @@
+use std::iter;
+
 use hir::Semantics;
 use ide_db::RootDatabase;
 use syntax::{
@@ -91,24 +93,42 @@ fn insert_whitespaces(syn: SyntaxNode) -> String {
         let is_last =
             |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };
 
-        res += &match token.kind() {
-            k if is_text(k) && is_next(|it| !it.is_punct(), true) => token.text().to_string() + " ",
+        match token.kind() {
+            k if is_text(k) && is_next(|it| !it.is_punct(), true) => {
+                res.push_str(token.text());
+                res.push(' ');
+            }
             L_CURLY if is_next(|it| it != R_CURLY, true) => {
                 indent += 1;
-                let leading_space = if is_last(is_text, false) { " " } else { "" };
-                format!("{}{{\n{}", leading_space, "  ".repeat(indent))
+                if is_last(is_text, false) {
+                    res.push(' ');
+                }
+                res.push_str("{\n");
+                res.extend(iter::repeat(" ").take(2 * indent));
             }
             R_CURLY if is_last(|it| it != L_CURLY, true) => {
                 indent = indent.saturating_sub(1);
-                format!("\n{}}}", "  ".repeat(indent))
+                res.push('\n');
+                res.extend(iter::repeat(" ").take(2 * indent));
+                res.push_str("}");
             }
-            R_CURLY => format!("}}\n{}", "  ".repeat(indent)),
-            T![;] => format!(";\n{}", "  ".repeat(indent)),
-            T![->] => " -> ".to_string(),
-            T![=] => " = ".to_string(),
-            T![=>] => " => ".to_string(),
-            _ => token.text().to_string(),
-        };
+            R_CURLY => {
+                res.push_str("}\n");
+                res.extend(iter::repeat(" ").take(2 * indent));
+            }
+            LIFETIME_IDENT if is_next(|it| it == IDENT, true) => {
+                res.push_str(token.text());
+                res.push(' ');
+            }
+            T![;] => {
+                res.push_str(";\n");
+                res.extend(iter::repeat(" ").take(2 * indent));
+            }
+            T![->] => res.push_str(" -> "),
+            T![=] => res.push_str(" = "),
+            T![=>] => res.push_str(" => "),
+            _ => res.push_str(token.text()),
+        }
 
         last = Some(token.kind());
     }
diff --git a/crates/ide/src/folding_ranges.rs b/crates/ide/src/folding_ranges.rs
index 153726ce8..2b9ed123c 100644
--- a/crates/ide/src/folding_ranges.rs
+++ b/crates/ide/src/folding_ranges.rs
@@ -19,6 +19,7 @@ pub enum FoldKind {
     Region,
     Consts,
     Statics,
+    Array,
 }
 
 #[derive(Debug)]
@@ -119,6 +120,7 @@ fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
     match kind {
         COMMENT => Some(FoldKind::Comment),
         ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
+        ARRAY_EXPR => Some(FoldKind::Array),
         ASSOC_ITEM_LIST
         | RECORD_FIELD_LIST
         | RECORD_PAT_FIELD_LIST
@@ -269,6 +271,7 @@ mod tests {
                 FoldKind::Region => "region",
                 FoldKind::Consts => "consts",
                 FoldKind::Statics => "statics",
+                FoldKind::Array => "array",
             };
             assert_eq!(kind, &attr.unwrap());
         }
@@ -465,6 +468,20 @@ fn foo<fold arglist>(
     }
 
     #[test]
+    fn fold_multiline_array() {
+        check(
+            r#"
+const FOO: [usize; 4] = <fold array>[
+    1,
+    2,
+    3,
+    4,
+]</fold>;
+"#,
+        )
+    }
+
+    #[test]
     fn fold_region() {
         check(
             r#"
diff --git a/crates/ide/src/syntax_highlighting/highlight.rs b/crates/ide/src/syntax_highlighting/highlight.rs
index 5ccb84714..e921784bf 100644
--- a/crates/ide/src/syntax_highlighting/highlight.rs
+++ b/crates/ide/src/syntax_highlighting/highlight.rs
@@ -323,8 +323,18 @@ fn highlight_def(db: &RootDatabase, def: Definition) -> Highlight {
             hir::ModuleDef::Trait(_) => HlTag::Symbol(SymbolKind::Trait),
             hir::ModuleDef::TypeAlias(type_) => {
                 let mut h = Highlight::new(HlTag::Symbol(SymbolKind::TypeAlias));
-                if type_.as_assoc_item(db).is_some() {
-                    h |= HlMod::Associated
+                if let Some(item) = type_.as_assoc_item(db) {
+                    h |= HlMod::Associated;
+                    match item.container(db) {
+                        AssocItemContainer::Impl(i) => {
+                            if i.trait_(db).is_some() {
+                                h |= HlMod::Trait;
+                            }
+                        }
+                        AssocItemContainer::Trait(_t) => {
+                            h |= HlMod::Trait;
+                        }
+                    }
                 }
                 return h;
             }
diff --git a/crates/ide/src/syntax_highlighting/tests.rs b/crates/ide/src/syntax_highlighting/tests.rs
index 1b02857ec..de2d22ac7 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -1,5 +1,8 @@
+use std::time::Instant;
+
 use expect_test::{expect_file, ExpectFile};
 use ide_db::SymbolKind;
+use stdx::format_to;
 use test_utils::{bench, bench_fixture, skip_slow_tests};
 
 use crate::{fixture, FileRange, HlTag, TextRange};
@@ -258,6 +261,99 @@ fn benchmark_syntax_highlighting_long_struct() {
 }
 
 #[test]
+fn syntax_highlighting_not_quadratic() {
+    if skip_slow_tests() {
+        return;
+    }
+
+    let mut measures = Vec::new();
+    for i in 6..=10 {
+        let n = 1 << i;
+        let fixture = bench_fixture::big_struct_n(n);
+        let (analysis, file_id) = fixture::file(&fixture);
+
+        let time = Instant::now();
+
+        let hash = analysis
+            .highlight(file_id)
+            .unwrap()
+            .iter()
+            .filter(|it| it.highlight.tag == HlTag::Symbol(SymbolKind::Struct))
+            .count();
+        assert!(hash > n as usize);
+
+        let elapsed = time.elapsed();
+        measures.push((n as f64, elapsed.as_millis() as f64))
+    }
+
+    assert_linear(&measures)
+}
+
+/// Checks that a set of measurements looks like a liner function rather than
+/// like a quadratic function. Algorithm:
+///
+/// 1. Linearly scale input to be in [0; 1)
+/// 2. Using linear regression, compute the best linear function approximating
+///    the input.
+/// 3. Compute RMSE and  maximal absolute error.
+/// 4. Check that errors are within tolerances and that the constant term is not
+///    too negative.
+///
+/// Ideally, we should use a proper "model selection" to directly compare
+/// quadratic and linear models, but that sounds rather complicated:
+///
+///     https://stats.stackexchange.com/questions/21844/selecting-best-model-based-on-linear-quadratic-and-cubic-fit-of-data
+fn assert_linear(xy: &[(f64, f64)]) {
+    let (mut xs, mut ys): (Vec<_>, Vec<_>) = xy.iter().copied().unzip();
+    normalize(&mut xs);
+    normalize(&mut ys);
+    let xy = xs.iter().copied().zip(ys.iter().copied());
+
+    // Linear regression: finding a and b to fit y = a + b*x.
+
+    let mean_x = mean(&xs);
+    let mean_y = mean(&ys);
+
+    let b = {
+        let mut num = 0.0;
+        let mut denom = 0.0;
+        for (x, y) in xy.clone() {
+            num += (x - mean_x) * (y - mean_y);
+            denom += (x - mean_x).powi(2);
+        }
+        num / denom
+    };
+
+    let a = mean_y - b * mean_x;
+
+    let mut plot = format!("y_pred = {:.3} + {:.3} * x\n\nx     y     y_pred\n", a, b);
+
+    let mut se = 0.0;
+    let mut max_error = 0.0f64;
+    for (x, y) in xy {
+        let y_pred = a + b * x;
+        se += (y - y_pred).powi(2);
+        max_error = max_error.max((y_pred - y).abs());
+
+        format_to!(plot, "{:.3} {:.3} {:.3}\n", x, y, y_pred);
+    }
+
+    let rmse = (se / xs.len() as f64).sqrt();
+    format_to!(plot, "\nrmse = {:.3} max error = {:.3}", rmse, max_error);
+
+    assert!(rmse < 0.05 && max_error < 0.1 && a > -0.1, "\nLooks quadratic\n{}", plot);
+
+    fn normalize(xs: &mut Vec<f64>) {
+        let max = xs.iter().copied().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap();
+        xs.iter_mut().for_each(|it| *it /= max);
+    }
+
+    fn mean(xs: &[f64]) -> f64 {
+        xs.iter().copied().sum::<f64>() / (xs.len() as f64)
+    }
+}
+
+#[test]
 fn benchmark_syntax_highlighting_parser() {
     if skip_slow_tests() {
         return;
diff --git a/crates/ide/src/typing.rs b/crates/ide/src/typing.rs
index 11408d445..1378048e5 100644
--- a/crates/ide/src/typing.rs
+++ b/crates/ide/src/typing.rs
@@ -22,18 +22,19 @@ use ide_db::{
 use syntax::{
     algo::find_node_at_offset,
     ast::{self, edit::IndentLevel, AstToken},
-    AstNode, SourceFile,
+    AstNode, Parse, SourceFile,
     SyntaxKind::{FIELD_EXPR, METHOD_CALL_EXPR},
     TextRange, TextSize,
 };
 
-use text_edit::TextEdit;
+use text_edit::{Indel, TextEdit};
 
 use crate::SourceChange;
 
 pub(crate) use on_enter::on_enter;
 
-pub(crate) const TRIGGER_CHARS: &str = ".=>";
+// Don't forget to add new trigger characters to `server_capabilities` in `caps.rs`.
+pub(crate) const TRIGGER_CHARS: &str = ".=>{";
 
 // Feature: On Typing Assists
 //
@@ -41,6 +42,7 @@ pub(crate) const TRIGGER_CHARS: &str = ".=>";
 //
 // - typing `let =` tries to smartly add `;` if `=` is followed by an existing expression
 // - typing `.` in a chain method call auto-indents
+// - typing `{` in front of an expression inserts a closing `}` after the expression
 //
 // VS Code::
 //
@@ -57,28 +59,79 @@ pub(crate) fn on_char_typed(
     position: FilePosition,
     char_typed: char,
 ) -> Option<SourceChange> {
-    assert!(TRIGGER_CHARS.contains(char_typed));
-    let file = &db.parse(position.file_id).tree();
-    assert_eq!(file.syntax().text().char_at(position.offset), Some(char_typed));
+    if !stdx::always!(TRIGGER_CHARS.contains(char_typed)) {
+        return None;
+    }
+    let file = &db.parse(position.file_id);
+    if !stdx::always!(file.tree().syntax().text().char_at(position.offset) == Some(char_typed)) {
+        return None;
+    }
     let edit = on_char_typed_inner(file, position.offset, char_typed)?;
     Some(SourceChange::from_text_edit(position.file_id, edit))
 }
 
-fn on_char_typed_inner(file: &SourceFile, offset: TextSize, char_typed: char) -> Option<TextEdit> {
-    assert!(TRIGGER_CHARS.contains(char_typed));
+fn on_char_typed_inner(
+    file: &Parse<SourceFile>,
+    offset: TextSize,
+    char_typed: char,
+) -> Option<TextEdit> {
+    if !stdx::always!(TRIGGER_CHARS.contains(char_typed)) {
+        return None;
+    }
     match char_typed {
-        '.' => on_dot_typed(file, offset),
-        '=' => on_eq_typed(file, offset),
-        '>' => on_arrow_typed(file, offset),
+        '.' => on_dot_typed(&file.tree(), offset),
+        '=' => on_eq_typed(&file.tree(), offset),
+        '>' => on_arrow_typed(&file.tree(), offset),
+        '{' => on_opening_brace_typed(file, offset),
         _ => unreachable!(),
     }
 }
 
+/// Inserts a closing `}` when the user types an opening `{`, wrapping an existing expression in a
+/// block.
+fn on_opening_brace_typed(file: &Parse<SourceFile>, offset: TextSize) -> Option<TextEdit> {
+    if !stdx::always!(file.tree().syntax().text().char_at(offset) == Some('{')) {
+        return None;
+    }
+
+    let brace_token = file.tree().syntax().token_at_offset(offset).right_biased()?;
+
+    // Remove the `{` to get a better parse tree, and reparse
+    let file = file.reparse(&Indel::delete(brace_token.text_range()));
+
+    let mut expr: ast::Expr = find_node_at_offset(file.tree().syntax(), offset)?;
+    if expr.syntax().text_range().start() != offset {
+        return None;
+    }
+
+    // Enclose the outermost expression starting at `offset`
+    while let Some(parent) = expr.syntax().parent() {
+        if parent.text_range().start() != expr.syntax().text_range().start() {
+            break;
+        }
+
+        match ast::Expr::cast(parent) {
+            Some(parent) => expr = parent,
+            None => break,
+        }
+    }
+
+    // If it's a statement in a block, we don't know how many statements should be included
+    if ast::ExprStmt::can_cast(expr.syntax().parent()?.kind()) {
+        return None;
+    }
+
+    // Insert `}` right after the expression.
+    Some(TextEdit::insert(expr.syntax().text_range().end() + TextSize::of("{"), "}".to_string()))
+}
+
 /// Returns an edit which should be applied after `=` was typed. Primarily,
 /// this works when adding `let =`.
 // FIXME: use a snippet completion instead of this hack here.
 fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
-    assert_eq!(file.syntax().text().char_at(offset), Some('='));
+    if !stdx::always!(file.syntax().text().char_at(offset) == Some('=')) {
+        return None;
+    }
     let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
     if let_stmt.semicolon_token().is_some() {
         return None;
@@ -100,7 +153,9 @@ fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
 
 /// Returns an edit which should be applied when a dot ('.') is typed on a blank line, indenting the line appropriately.
 fn on_dot_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
-    assert_eq!(file.syntax().text().char_at(offset), Some('.'));
+    if !stdx::always!(file.syntax().text().char_at(offset) == Some('.')) {
+        return None;
+    }
     let whitespace =
         file.syntax().token_at_offset(offset).left_biased().and_then(ast::Whitespace::cast)?;
 
@@ -129,7 +184,9 @@ fn on_dot_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
 /// Adds a space after an arrow when `fn foo() { ... }` is turned into `fn foo() -> { ... }`
 fn on_arrow_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
     let file_text = file.syntax().text();
-    assert_eq!(file_text.char_at(offset), Some('>'));
+    if !stdx::always!(file_text.char_at(offset) == Some('>')) {
+        return None;
+    }
     let after_arrow = offset + TextSize::of('>');
     if file_text.char_at(after_arrow) != Some('{') {
         return None;
@@ -152,7 +209,7 @@ mod tests {
         let edit = TextEdit::insert(offset, char_typed.to_string());
         edit.apply(&mut before);
         let parse = SourceFile::parse(&before);
-        on_char_typed_inner(&parse.tree(), offset, char_typed).map(|it| {
+        on_char_typed_inner(&parse, offset, char_typed).map(|it| {
             it.apply(&mut before);
             before.to_string()
         })
@@ -373,4 +430,85 @@ fn main() {
     fn adds_space_after_return_type() {
         type_char('>', "fn foo() -$0{ 92 }", "fn foo() -> { 92 }")
     }
+
+    #[test]
+    fn adds_closing_brace() {
+        type_char(
+            '{',
+            r#"
+fn f() { match () { _ => $0() } }
+            "#,
+            r#"
+fn f() { match () { _ => {()} } }
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+fn f() { $0() }
+            "#,
+            r#"
+fn f() { {()} }
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+fn f() { let x = $0(); }
+            "#,
+            r#"
+fn f() { let x = {()}; }
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+fn f() { let x = $0a.b(); }
+            "#,
+            r#"
+fn f() { let x = {a.b()}; }
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+const S: () = $0();
+fn f() {}
+            "#,
+            r#"
+const S: () = {()};
+fn f() {}
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+const S: () = $0a.b();
+fn f() {}
+            "#,
+            r#"
+const S: () = {a.b()};
+fn f() {}
+            "#,
+        );
+        type_char(
+            '{',
+            r#"
+fn f() {
+    match x {
+        0 => $0(),
+        1 => (),
+    }
+}
+            "#,
+            r#"
+fn f() {
+    match x {
+        0 => {()},
+        1 => (),
+    }
+}
+            "#,
+        );
+    }
 }
diff --git a/crates/ide/src/typing/on_enter.rs b/crates/ide/src/typing/on_enter.rs
index 9144681bf..6f1ce3689 100644
--- a/crates/ide/src/typing/on_enter.rs
+++ b/crates/ide/src/typing/on_enter.rs
@@ -18,6 +18,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 `//!`
 //
 // This action needs to be assigned to shortcut explicitly.
 //
@@ -187,6 +188,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"