6 files changed, 187 insertions, 79 deletions

diff --git a/crates/ra_editor/src/folding_ranges.rs b/crates/ra_editor/src/folding_ranges.rs
index 3aabd54ae..a1699d449 100644
--- a/crates/ra_editor/src/folding_ranges.rs
+++ b/crates/ra_editor/src/folding_ranges.rs
@@ -1,8 +1,10 @@
 use rustc_hash::FxHashSet;
 
 use ra_syntax::{
+    ast,
+    AstNode,
     File, TextRange, SyntaxNodeRef,
-    SyntaxKind,
+    SyntaxKind::{self, *},
     Direction,
 };
 
@@ -20,67 +22,97 @@ pub struct Fold {
 
 pub fn folding_ranges(file: &File) -> Vec<Fold> {
     let mut res = vec![];
-    let mut visited = FxHashSet::default();
+    let mut visited_comments = FxHashSet::default();
 
     for node in file.syntax().descendants() {
-        if visited.contains(&node) {
+        // Fold items that span multiple lines
+        if let Some(kind) = fold_kind(node.kind()) {
+            if has_newline(node) {
+                res.push(Fold { range: node.range(), kind });
+            }
+        }
+
+        // Also fold groups of comments
+        if visited_comments.contains(&node) {
             continue;
         }
+        if node.kind() == COMMENT {
+            contiguous_range_for_comment(node, &mut visited_comments)
+                .map(|range| res.push(Fold { range, kind: FoldKind::Comment }));
+        }
+    }
+
+    res
+}
 
-        let range_and_kind = match node.kind() {
-            SyntaxKind::COMMENT => (
-                contiguous_range_for(SyntaxKind::COMMENT, node, &mut visited),
-                Some(FoldKind::Comment),
-            ),
-            SyntaxKind::USE_ITEM => (
-                contiguous_range_for(SyntaxKind::USE_ITEM, node, &mut visited),
-                Some(FoldKind::Imports),
-            ),
-            _ => (None, None),
-        };
-
-        match range_and_kind {
-            (Some(range), Some(kind)) => {
-                res.push(Fold {
-                    range: range,
-                    kind: kind
-                });
+fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
+    match kind {
+        COMMENT => Some(FoldKind::Comment),
+        USE_ITEM => Some(FoldKind::Imports),
+        _ => None
+    }
+}
+
+fn has_newline(
+    node: SyntaxNodeRef,
+) -> bool {
+    for descendant in node.descendants() {
+        if let Some(ws) = ast::Whitespace::cast(descendant) {
+            if ws.has_newlines() {
+                return true;
+            }
+        } else if let Some(comment) = ast::Comment::cast(descendant) {
+            if comment.has_newlines() {
+                return true;
             }
-            _ => {}
         }
     }
 
-    res
+    false
 }
 
-fn contiguous_range_for<'a>(
-    kind: SyntaxKind,
-    node: SyntaxNodeRef<'a>,
+fn contiguous_range_for_comment<'a>(
+    first: SyntaxNodeRef<'a>,
     visited: &mut FxHashSet<SyntaxNodeRef<'a>>,
 ) -> Option<TextRange> {
-    visited.insert(node);
-
-    let left = node;
-    let mut right = node;
-    for node in node.siblings(Direction::Next) {
-        visited.insert(node);
-        match node.kind() {
-            SyntaxKind::WHITESPACE if !node.leaf_text().unwrap().as_str().contains("\n\n") => (),
-            k => {
-                if k == kind {
-                    right = node
-                } else {
-                    break;
-                }
+    visited.insert(first);
+
+    // Only fold comments of the same flavor
+    let group_flavor = ast::Comment::cast(first)?.flavor();
+
+    let mut last = first;
+    for node in first.siblings(Direction::Next) {
+        if let Some(ws) = ast::Whitespace::cast(node) {
+            // There is a blank line, which means the group ends here
+            if ws.count_newlines_lazy().take(2).count() == 2 {
+                break;
+            }
+
+            // Ignore whitespace without blank lines
+            continue;
+        }
+
+        match ast::Comment::cast(node) {
+            Some(next_comment) if next_comment.flavor() == group_flavor => {
+                visited.insert(node);
+                last = node;
+            }
+            // The comment group ends because either:
+            // * An element of a different kind was reached
+            // * A comment of a different flavor was reached
+            _ => {
+                break
             }
         }
     }
-    if left != right {
+
+    if first != last {
         Some(TextRange::from_to(
-            left.range().start(),
-            right.range().end(),
+            first.range().start(),
+            last.range().end(),
         ))
     } else {
+        // The group consists of only one element, therefore it cannot be folded
         None
     }
 }
@@ -88,52 +120,65 @@ fn contiguous_range_for<'a>(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use test_utils::extract_ranges;
+
+    fn do_check(text: &str, fold_kinds: &[FoldKind]) {
+        let (ranges, text) = extract_ranges(text);
+        let file = File::parse(&text);
+        let folds = folding_ranges(&file);
+
+        assert_eq!(folds.len(), ranges.len());
+        for ((fold, range), fold_kind) in folds.into_iter().zip(ranges.into_iter()).zip(fold_kinds.into_iter()) {
+            assert_eq!(fold.range.start(), range.start());
+            assert_eq!(fold.range.end(), range.end());
+            assert_eq!(&fold.kind, fold_kind);
+        }
+    }
 
     #[test]
     fn test_fold_comments() {
         let text = r#"
-// Hello
+<|>// Hello
 // this is a multiline
 // comment
-//
+//<|>
 
 // But this is not
 
 fn main() {
-    // We should
+    <|>// We should
     // also
     // fold
-    // this one.
+    // this one.<|>
+    <|>//! But this one is different
+    //! because it has another flavor<|>
+    <|>/* As does this
+    multiline comment */<|>
 }"#;
 
-        let file = File::parse(&text);
-        let folds = folding_ranges(&file);
-        assert_eq!(folds.len(), 2);
-        assert_eq!(folds[0].range.start(), 1.into());
-        assert_eq!(folds[0].range.end(), 46.into());
-        assert_eq!(folds[0].kind, FoldKind::Comment);
-
-        assert_eq!(folds[1].range.start(), 84.into());
-        assert_eq!(folds[1].range.end(), 137.into());
-        assert_eq!(folds[1].kind, FoldKind::Comment);
+        let fold_kinds = &[
+            FoldKind::Comment,
+            FoldKind::Comment,
+            FoldKind::Comment,
+            FoldKind::Comment,
+        ];
+        do_check(text, fold_kinds);
     }
 
     #[test]
     fn test_fold_imports() {
         let text = r#"
-use std::str;
-use std::vec;
-use std::io as iop;
+<|>use std::{
+    str,
+    vec,
+    io as iop
+};<|>
 
 fn main() {
 }"#;
 
-        let file = File::parse(&text);
-        let folds = folding_ranges(&file);
-        assert_eq!(folds.len(), 1);
-        assert_eq!(folds[0].range.start(), 1.into());
-        assert_eq!(folds[0].range.end(), 48.into());
-        assert_eq!(folds[0].kind, FoldKind::Imports);
+        let folds = &[FoldKind::Imports];
+        do_check(text, folds);
     }
 
 
diff --git a/crates/ra_syntax/src/ast/generated.rs b/crates/ra_syntax/src/ast/generated.rs
index 48c9038dc..4db1bcbf9 100644
--- a/crates/ra_syntax/src/ast/generated.rs
+++ b/crates/ra_syntax/src/ast/generated.rs
@@ -2197,3 +2197,21 @@ impl<'a> WhileExpr<'a> {
     }
 }
 
+// Whitespace
+#[derive(Debug, Clone, Copy)]
+pub struct Whitespace<'a> {
+    syntax: SyntaxNodeRef<'a>,
+}
+
+impl<'a> AstNode<'a> for Whitespace<'a> {
+    fn cast(syntax: SyntaxNodeRef<'a>) -> Option<Self> {
+        match syntax.kind() {
+            WHITESPACE => Some(Whitespace { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(self) -> SyntaxNodeRef<'a> { self.syntax }
+}
+
+impl<'a> Whitespace<'a> {}
+
diff --git a/crates/ra_syntax/src/ast/mod.rs b/crates/ra_syntax/src/ast/mod.rs
index 10dac72e5..12ddc0210 100644
--- a/crates/ra_syntax/src/ast/mod.rs
+++ b/crates/ra_syntax/src/ast/mod.rs
@@ -100,8 +100,8 @@ impl<'a> Lifetime<'a> {
 }
 
 impl<'a> Comment<'a> {
-    pub fn text(&self) -> SmolStr {
-        self.syntax().leaf_text().unwrap().clone()
+    pub fn text(&self) -> &SmolStr {
+        self.syntax().leaf_text().unwrap()
     }
 
     pub fn flavor(&self) -> CommentFlavor {
@@ -120,9 +120,17 @@ impl<'a> Comment<'a> {
     pub fn prefix(&self) -> &'static str {
         self.flavor().prefix()
     }
+
+    pub fn count_newlines_lazy(&self) -> impl Iterator<Item = &()> {
+        self.text().chars().filter(|&c| c == '\n').map(|_| &())
+    }
+
+    pub fn has_newlines(&self) -> bool {
+        self.count_newlines_lazy().count() > 0
+    }
 }
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum CommentFlavor {
     Line,
     Doc,
@@ -142,6 +150,20 @@ impl CommentFlavor {
     }
 }
 
+impl<'a> Whitespace<'a> {
+    pub fn text(&self) -> &SmolStr {
+        &self.syntax().leaf_text().unwrap()
+    }
+
+    pub fn count_newlines_lazy(&self) -> impl Iterator<Item = &()> {
+        self.text().chars().filter(|&c| c == '\n').map(|_| &())
+    }
+
+    pub fn has_newlines(&self) -> bool {
+        self.count_newlines_lazy().count() > 0
+    }
+}
+
 impl<'a> Name<'a> {
     pub fn text(&self) -> SmolStr {
         let ident = self.syntax().first_child()
diff --git a/crates/ra_syntax/src/grammar.ron b/crates/ra_syntax/src/grammar.ron
index a904f7505..ea8063d3b 100644
--- a/crates/ra_syntax/src/grammar.ron
+++ b/crates/ra_syntax/src/grammar.ron
@@ -538,5 +538,6 @@ Grammar(
             options: [ "NameRef" ]
         ),
         "Comment": (),
+        "Whitespace": (),
     },
 )
diff --git a/crates/test_utils/src/lib.rs b/crates/test_utils/src/lib.rs
index 068eb80ce..ee73153f0 100644
--- a/crates/test_utils/src/lib.rs
+++ b/crates/test_utils/src/lib.rs
@@ -38,22 +38,44 @@ pub fn assert_eq_dbg(expected: &str, actual: &impl fmt::Debug) {
 }
 
 pub fn extract_offset(text: &str) -> (TextUnit, String) {
-    let cursor = "<|>";
-    let cursor_pos = match text.find(cursor) {
+    match try_extract_offset(text) {
         None => panic!("text should contain cursor marker"),
-        Some(pos) => pos,
-    };
+        Some(result) => result,
+    }
+}
+
+pub fn try_extract_offset(text: &str) -> Option<(TextUnit, String)> {
+    let cursor = "<|>";
+    let cursor_pos = text.find(cursor)?;
     let mut new_text = String::with_capacity(text.len() - cursor.len());
     new_text.push_str(&text[..cursor_pos]);
     new_text.push_str(&text[cursor_pos + cursor.len()..]);
     let cursor_pos = TextUnit::from(cursor_pos as u32);
-    (cursor_pos, new_text)
+    Some((cursor_pos, new_text))
 }
 
 pub fn extract_range(text: &str) -> (TextRange, String) {
-    let (start, text) = extract_offset(text);
-    let (end, text) = extract_offset(&text);
-    (TextRange::from_to(start, end), text)
+    match try_extract_range(text) {
+        None => panic!("text should contain cursor marker"),
+        Some(result) => result,
+    }
+}
+
+pub fn try_extract_range(text: &str) -> Option<(TextRange, String)> {
+    let (start, text) = try_extract_offset(text)?;
+    let (end, text) = try_extract_offset(&text)?;
+    Some((TextRange::from_to(start, end), text))
+}
+
+pub fn extract_ranges(text: &str) -> (Vec<TextRange>, String) {
+    let mut ranges = Vec::new();
+    let mut text = String::from(text);
+    while let Some((range, new_text)) = try_extract_range(&text) {
+        text = new_text;
+        ranges.push(range);
+    }
+
+    (ranges, text)
 }
 
 pub fn add_cursor(text: &str, offset: TextUnit) -> String {
diff --git a/editors/code/src/extension.ts b/editors/code/src/extension.ts
index ff8f23c7a..d1c525f68 100644
--- a/editors/code/src/extension.ts
+++ b/editors/code/src/extension.ts
@@ -20,8 +20,8 @@ export function activate(context: vscode.ExtensionContext) {
         f: (...args: any[]) => Promise<boolean>
     ) {
         const defaultCmd = `default:${name}`;
-        const original = async (...args: any[]) =>
-            await vscode.commands.executeCommand(defaultCmd, ...args);
+        const original = (...args: any[]) =>
+            vscode.commands.executeCommand(defaultCmd, ...args);
 
         registerCommand(name, async (...args: any[]) => {
             const editor = vscode.window.activeTextEditor;