Cleanup fold code and split logic to fold single elements

4 files changed, 112 insertions, 40 deletions

diff --git a/crates/ra_editor/src/folding_ranges.rs b/crates/ra_editor/src/folding_ranges.rs
index 1d89ac853..d45210ec7 100644
--- a/crates/ra_editor/src/folding_ranges.rs
+++ b/crates/ra_editor/src/folding_ranges.rs
@@ -22,65 +22,96 @@ pub struct Fold {
 
 pub fn folding_ranges(file: &File) -> Vec<Fold> {
     let mut res = vec![];
-    let mut visited = FxHashSet::default();
+    let mut group_members = FxHashSet::default();
 
     for node in file.syntax().descendants() {
-        if visited.contains(&node) {
-            continue;
+        // 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 });
+            }
         }
 
-        if let Some(comment) = ast::Comment::cast(node) {
-            // Multiline comments (`/* ... */`) can only be folded if they span multiple lines
-            let range = if let ast::CommentFlavor::Multiline = comment.flavor() {
-                if comment.text().contains('\n') {
-                    Some(comment.syntax().range())
-                } else {
-                    None
-                }
-            } else {
-                contiguous_range_for(SyntaxKind::COMMENT, node, &mut visited)
-            };
-
-            range.map(|range| res.push(Fold { range, kind: FoldKind::Comment }));
+        // Also fold item *groups* that span multiple lines
+
+        // Note: we need to skip elements of the group that we have already visited,
+        // otherwise there will be folds for the whole group and for its sub groups
+        if group_members.contains(&node) {
+            continue;
         }
 
-        if let SyntaxKind::USE_ITEM = node.kind() {
-            contiguous_range_for(SyntaxKind::USE_ITEM, node, &mut visited)
-                .map(|range| res.push(Fold { range, kind: FoldKind::Imports}));
-        };
+        if let Some(kind) = fold_kind(node.kind()) {
+            contiguous_range_for_group(node.kind(), node, &mut group_members)
+                .map(|range| res.push(Fold { range, kind }));
+        }
     }
 
     res
 }
 
-fn contiguous_range_for<'a>(
-    kind: SyntaxKind,
-    node: SyntaxNodeRef<'a>,
+fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
+    match kind {
+        SyntaxKind::COMMENT => Some(FoldKind::Comment),
+        SyntaxKind::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;
+            }
+        }
+    }
+
+    false
+}
+
+fn contiguous_range_for_group<'a>(
+    group_kind: SyntaxKind,
+    first: SyntaxNodeRef<'a>,
     visited: &mut FxHashSet<SyntaxNodeRef<'a>>,
 ) -> Option<TextRange> {
-    visited.insert(node);
+    visited.insert(first);
+
+    let mut last = first;
 
-    let left = node;
-    let mut right = node;
-    for node in node.siblings(Direction::Next) {
+    for node in first.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;
-                }
+        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;
+        }
+
+        // The group ends when an element of a different kind is reached
+        if node.kind() != group_kind {
+            break;
         }
+
+        // Keep track of the last node in the group
+        last = node;
     }
-    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
     }
 }
diff --git a/crates/ra_syntax/src/ast/generated.rs b/crates/ra_syntax/src/ast/generated.rs
index ef7b5b1a1..85aa5e0dd 100644
--- a/crates/ra_syntax/src/ast/generated.rs
+++ b/crates/ra_syntax/src/ast/generated.rs
@@ -2193,3 +2193,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..00c852274 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,6 +120,14 @@ 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)]
@@ -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 9da0c2c13..d538739de 100644
--- a/crates/ra_syntax/src/grammar.ron
+++ b/crates/ra_syntax/src/grammar.ron
@@ -538,5 +538,6 @@ Grammar(
             options: [ "NameRef" ]
         ),
         "Comment": (),
+        "Whitespace": (),
     },
 )