6 files changed, 184 insertions, 22 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 d5628e3df..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,27 +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 => {
+                res.push_str("}\n");
+                res.extend(iter::repeat(" ").take(2 * indent));
             }
-            R_CURLY => format!("}}\n{}", "  ".repeat(indent)),
             LIFETIME_IDENT if is_next(|it| it == IDENT, true) => {
-                format!("{} ", token.text().to_string())
+                res.push_str(token.text());
+                res.push(' ');
             }
-            T![;] => format!(";\n{}", "  ".repeat(indent)),
-            T![->] => " -> ".to_string(),
-            T![=] => " = ".to_string(),
-            T![=>] => " => ".to_string(),
-            _ => token.text().to_string(),
-        };
+            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/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"