Merge pull request #20080 from Veykril/push-vnrwqppplykm

Cleanup `folding_ranges` and support more things

2 files changed, 60 insertions, 39 deletions

diff --git a/src/tools/rust-analyzer/crates/ide/src/folding_ranges.rs b/src/tools/rust-analyzer/crates/ide/src/folding_ranges.rs
index 194e8c968f7..9bd8504733a 100755
--- a/src/tools/rust-analyzer/crates/ide/src/folding_ranges.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/folding_ranges.rs
@@ -2,7 +2,7 @@ use ide_db::{FxHashSet, syntax_helpers::node_ext::vis_eq};
 use syntax::{
     Direction, NodeOrToken, SourceFile,
     SyntaxKind::{self, *},
-    TextRange, TextSize,
+    SyntaxNode, TextRange, TextSize,
     ast::{self, AstNode, AstToken},
     match_ast,
 };
@@ -16,16 +16,21 @@ const REGION_END: &str = "// endregion";
 pub enum FoldKind {
     Comment,
     Imports,
-    Mods,
+    Region,
     Block,
     ArgList,
-    Region,
-    Consts,
-    Statics,
     Array,
     WhereClause,
     ReturnType,
     MatchArm,
+    // region: item runs
+    Modules,
+    Consts,
+    Statics,
+    TypeAliases,
+    TraitAliases,
+    ExternCrates,
+    // endregion: item runs
 }
 
 #[derive(Debug)]
@@ -41,10 +46,7 @@ pub struct Fold {
 pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
     let mut res = vec![];
     let mut visited_comments = FxHashSet::default();
-    let mut visited_imports = FxHashSet::default();
-    let mut visited_mods = FxHashSet::default();
-    let mut visited_consts = FxHashSet::default();
-    let mut visited_statics = FxHashSet::default();
+    let mut visited_nodes = FxHashSet::default();
 
     // regions can be nested, here is a LIFO buffer
     let mut region_starts: Vec<TextSize> = vec![];
@@ -93,30 +95,40 @@ pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
                             if module.item_list().is_none() {
                                 if let Some(range) = contiguous_range_for_item_group(
                                     module,
-                                    &mut visited_mods,
+                                    &mut visited_nodes,
                                 ) {
-                                    res.push(Fold { range, kind: FoldKind::Mods })
+                                    res.push(Fold { range, kind: FoldKind::Modules })
                                 }
                             }
                         },
                         ast::Use(use_) => {
-                            if let Some(range) = contiguous_range_for_item_group(use_, &mut visited_imports) {
+                            if let Some(range) = contiguous_range_for_item_group(use_, &mut visited_nodes) {
                                 res.push(Fold { range, kind: FoldKind::Imports })
                             }
                         },
                         ast::Const(konst) => {
-                            if let Some(range) = contiguous_range_for_item_group(konst, &mut visited_consts) {
+                            if let Some(range) = contiguous_range_for_item_group(konst, &mut visited_nodes) {
                                 res.push(Fold { range, kind: FoldKind::Consts })
                             }
                         },
                         ast::Static(statik) => {
-                            if let Some(range) = contiguous_range_for_item_group(statik, &mut visited_statics) {
+                            if let Some(range) = contiguous_range_for_item_group(statik, &mut visited_nodes) {
                                 res.push(Fold { range, kind: FoldKind::Statics })
                             }
                         },
-                        ast::WhereClause(where_clause) => {
-                            if let Some(range) = fold_range_for_where_clause(where_clause) {
-                                res.push(Fold { range, kind: FoldKind::WhereClause })
+                        ast::TypeAlias(alias) => {
+                            if let Some(range) = contiguous_range_for_item_group(alias, &mut visited_nodes) {
+                                res.push(Fold { range, kind: FoldKind::TypeAliases })
+                            }
+                        },
+                        ast::TraitAlias(alias) => {
+                            if let Some(range) = contiguous_range_for_item_group(alias, &mut visited_nodes) {
+                                res.push(Fold { range, kind: FoldKind::TraitAliases })
+                            }
+                        },
+                        ast::ExternCrate(extern_crate) => {
+                            if let Some(range) = contiguous_range_for_item_group(extern_crate, &mut visited_nodes) {
+                                res.push(Fold { range, kind: FoldKind::ExternCrates })
                             }
                         },
                         ast::MatchArm(match_arm) => {
@@ -137,9 +149,10 @@ pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
 fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
     match kind {
         COMMENT => Some(FoldKind::Comment),
-        ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
+        ARG_LIST | PARAM_LIST | GENERIC_ARG_LIST | GENERIC_PARAM_LIST => Some(FoldKind::ArgList),
         ARRAY_EXPR => Some(FoldKind::Array),
         RET_TYPE => Some(FoldKind::ReturnType),
+        WHERE_CLAUSE => Some(FoldKind::WhereClause),
         ASSOC_ITEM_LIST
         | RECORD_FIELD_LIST
         | RECORD_PAT_FIELD_LIST
@@ -155,11 +168,14 @@ fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
     }
 }
 
-fn contiguous_range_for_item_group<N>(first: N, visited: &mut FxHashSet<N>) -> Option<TextRange>
+fn contiguous_range_for_item_group<N>(
+    first: N,
+    visited: &mut FxHashSet<SyntaxNode>,
+) -> Option<TextRange>
 where
     N: ast::HasVisibility + Clone + Hash + Eq,
 {
-    if !visited.insert(first.clone()) {
+    if !visited.insert(first.syntax().clone()) {
         return None;
     }
 
@@ -183,7 +199,7 @@ where
         if let Some(next) = N::cast(node) {
             let next_vis = next.visibility();
             if eq_visibility(next_vis.clone(), last_vis) {
-                visited.insert(next.clone());
+                visited.insert(next.syntax().clone());
                 last_vis = next_vis;
                 last = next;
                 continue;
@@ -259,18 +275,6 @@ fn contiguous_range_for_comment(
     }
 }
 
-fn fold_range_for_where_clause(where_clause: ast::WhereClause) -> Option<TextRange> {
-    let first_where_pred = where_clause.predicates().next();
-    let last_where_pred = where_clause.predicates().last();
-
-    if first_where_pred != last_where_pred {
-        let start = where_clause.where_token()?.text_range().end();
-        let end = where_clause.syntax().text_range().end();
-        return Some(TextRange::new(start, end));
-    }
-    None
-}
-
 fn fold_range_for_multiline_match_arm(match_arm: ast::MatchArm) -> Option<TextRange> {
     if fold_kind(match_arm.expr()?.syntax().kind()).is_some() {
         None
@@ -307,16 +311,19 @@ mod tests {
             let kind = match fold.kind {
                 FoldKind::Comment => "comment",
                 FoldKind::Imports => "imports",
-                FoldKind::Mods => "mods",
+                FoldKind::Modules => "mods",
                 FoldKind::Block => "block",
                 FoldKind::ArgList => "arglist",
                 FoldKind::Region => "region",
                 FoldKind::Consts => "consts",
                 FoldKind::Statics => "statics",
+                FoldKind::TypeAliases => "typealiases",
                 FoldKind::Array => "array",
                 FoldKind::WhereClause => "whereclause",
                 FoldKind::ReturnType => "returntype",
                 FoldKind::MatchArm => "matcharm",
+                FoldKind::TraitAliases => "traitaliases",
+                FoldKind::ExternCrates => "externcrates",
             };
             assert_eq!(kind, &attr.unwrap());
         }
@@ -594,19 +601,18 @@ static SECOND_STATIC: &str = "second";</fold>
 
     #[test]
     fn fold_where_clause() {
-        // fold multi-line and don't fold single line.
         check(
             r#"
 fn foo()
-where<fold whereclause>
+<fold whereclause>where
     A: Foo,
     B: Foo,
     C: Foo,
     D: Foo,</fold> {}
 
 fn bar()
-where
-    A: Bar, {}
+<fold whereclause>where
+    A: Bar,</fold> {}
 "#,
         )
     }
@@ -624,4 +630,16 @@ fn bar() -> (bool, bool) { (true, true) }
 "#,
         )
     }
+
+    #[test]
+    fn fold_generics() {
+        check(
+            r#"
+type Foo<T, U> = foo<fold arglist><
+    T,
+    U,
+></fold>;
+"#,
+        )
+    }
 }
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
index 4efe330f16a..8a848fb848c 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
@@ -900,14 +900,17 @@ pub(crate) fn folding_range(
         FoldKind::Comment => Some(lsp_types::FoldingRangeKind::Comment),
         FoldKind::Imports => Some(lsp_types::FoldingRangeKind::Imports),
         FoldKind::Region => Some(lsp_types::FoldingRangeKind::Region),
-        FoldKind::Mods
+        FoldKind::Modules
         | FoldKind::Block
         | FoldKind::ArgList
         | FoldKind::Consts
         | FoldKind::Statics
+        | FoldKind::TypeAliases
         | FoldKind::WhereClause
         | FoldKind::ReturnType
         | FoldKind::Array
+        | FoldKind::TraitAliases
+        | FoldKind::ExternCrates
         | FoldKind::MatchArm => None,
     };