✨ feat(mq-lsp): add folding range support

crates/mq-lsp/README.md

@@ -13,6 +13,7 @@ Language Server Protocol (LSP) implementation for the [mq](https://mqlang.org/)
 - 🔗 **Find References**: Locate all usages of a symbol across your workspace
 - 🗂️ **Document Symbols**: Outline view of all symbols in the current file
 - 🗃️ **Workspace Symbols**: Search for symbols by name across all loaded files/modules
+- 📁 **Folding Ranges**: Collapse function/macro/module bodies, control-flow blocks, multi-line array/dict literals, and multi-line comment banners
 - 🎨 **Semantic Tokens**: Enhanced syntax highlighting based on semantic analysis
 - ✨ **Code Formatting**: Automatic code formatting following mq style guidelines
 - 🛠️ **Code Actions**: Quick fixes such as adding a missing `include`/`import` for an unresolved function or module reference

crates/mq-lsp/src/capabilities.rs

@@ -1,9 +1,10 @@
 use tower_lsp_server::ls_types::{
     CodeActionKind, CodeActionOptions, CodeActionProviderCapability, CompletionOptions, DiagnosticOptions,
     DiagnosticServerCapabilities, DocumentFormattingOptions, DocumentRangeFormattingOptions, ExecuteCommandOptions,
-    HoverProviderCapability, InlayHintOptions, InlayHintServerCapabilities, OneOf, RenameOptions,
-    SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions, SemanticTokensServerCapabilities,
-    ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
+    FoldingRangeProviderCapability, HoverProviderCapability, InlayHintOptions, InlayHintServerCapabilities, OneOf,
+    RenameOptions, SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions,
+    SemanticTokensServerCapabilities, ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability,
+    TextDocumentSyncKind,
 };
 
 use crate::semantic_tokens;
@@ -41,6 +42,7 @@ pub(crate) fn server_capabilities() -> ServerCapabilities {
         })),
         document_symbol_provider: Some(OneOf::Left(true)),
         workspace_symbol_provider: Some(OneOf::Left(true)),
+        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
         diagnostic_provider: Some(DiagnosticServerCapabilities::Options(DiagnosticOptions {
             identifier: None,
             // Editing a module/import can change diagnostics in files that depend on it.

crates/mq-lsp/src/folding_range.rs

@@ -0,0 +1,173 @@
+use mq_lang::{CstNode, CstNodeKind, CstTrivia, Shared};
+use tower_lsp_server::ls_types::{FoldingRange, FoldingRangeKind};
+
+/// Computes folding ranges by parsing the document into a CST and walking it directly,
+/// the same approach `signature_help.rs` uses — `mq_hir::Symbol` ranges are too coarse
+/// (block delimiters like `end`/`;` aren't lowered into HIR at all), but the CST's
+/// `node_range()` gives an exact span for every block construct.
+pub(crate) fn response(source_text: Option<&str>) -> Option<Vec<FoldingRange>> {
+    let source_text = source_text?;
+    let (nodes, _) = mq_lang::parse_recovery(source_text);
+
+    let mut ranges = Vec::new();
+    for node in &nodes {
+        visit(node, &mut ranges);
+    }
+
+    if ranges.is_empty() { None } else { Some(ranges) }
+}
+
+fn visit(node: &Shared<CstNode>, ranges: &mut Vec<FoldingRange>) {
+    collect_comment_folds(node, ranges);
+
+    if is_foldable(&node.kind) {
+        let span = node.node_range();
+        if span.end.line > span.start.line {
+            ranges.push(FoldingRange {
+                start_line: span.start.line - 1,
+                start_character: None,
+                end_line: span.end.line - 1,
+                end_character: None,
+                kind: Some(FoldingRangeKind::Region),
+                collapsed_text: None,
+            });
+        }
+    }
+
+    for child in &node.children {
+        visit(child, ranges);
+    }
+}
+
+/// Block-like constructs worth collapsing: function/macro/module bodies, control-flow
+/// blocks, and multi-line array/dict literals. Deliberately excludes leaf/expression kinds
+/// (`Call`, `BinaryOp`, ...) so folding stays limited to structural blocks.
+fn is_foldable(kind: &CstNodeKind) -> bool {
+    matches!(
+        kind,
+        CstNodeKind::Def
+            | CstNodeKind::Macro
+            | CstNodeKind::Module
+            | CstNodeKind::If
+            | CstNodeKind::Elif
+            | CstNodeKind::Else
+            | CstNodeKind::Match
+            | CstNodeKind::MatchArm
+            | CstNodeKind::Foreach
+            | CstNodeKind::While
+            | CstNodeKind::Loop
+            | CstNodeKind::Try
+            | CstNodeKind::Catch
+            | CstNodeKind::Array
+            | CstNodeKind::Dict
+    )
+}
+
+/// Folds runs of 2+ consecutive `#`-comment lines immediately preceding a node, so a
+/// multi-line doc comment or a `# Section:`-style banner can be collapsed like in other
+/// languages' "region"/"comment" folding.
+fn collect_comment_folds(node: &Shared<CstNode>, ranges: &mut Vec<FoldingRange>) {
+    let comment_lines = node
+        .leading_trivia
+        .iter()
+        .filter_map(|trivia| match trivia {
+            CstTrivia::Comment(token) => Some(token.range.start.line),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+
+    let mut run_start = None;
+    let mut run_end = None;
+
+    for line in comment_lines {
+        match run_end {
+            Some(end) if line == end + 1 => run_end = Some(line),
+            _ => {
+                push_comment_fold(ranges, run_start, run_end);
+                run_start = Some(line);
+                run_end = Some(line);
+            }
+        }
+    }
+    push_comment_fold(ranges, run_start, run_end);
+}
+
+fn push_comment_fold(ranges: &mut Vec<FoldingRange>, start: Option<u32>, end: Option<u32>) {
+    if let (Some(start), Some(end)) = (start, end)
+        && end > start
+    {
+        ranges.push(FoldingRange {
+            start_line: start - 1,
+            start_character: None,
+            end_line: end - 1,
+            end_character: None,
+            kind: Some(FoldingRangeKind::Comment),
+            collapsed_text: None,
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_no_source_text() {
+        assert!(response(None).is_none());
+    }
+
+    #[test]
+    fn test_no_foldable_ranges_for_single_line() {
+        let result = response(Some("let x = 1"));
+        assert!(result.is_none());
+    }
+
+    #[test]
+    fn test_folds_multiline_function_body() {
+        let code = "def foo(a):\n  let b = a + 1\n  | b;\n| foo(1)";
+        let result = response(Some(code)).unwrap();
+
+        let def_fold = result
+            .iter()
+            .find(|r| r.kind == Some(FoldingRangeKind::Region) && r.start_line == 0);
+        assert!(
+            def_fold.is_some(),
+            "expected a folding range starting at the `def` line"
+        );
+        assert_eq!(def_fold.unwrap().end_line, 2);
+    }
+
+    #[test]
+    fn test_folds_multiline_array() {
+        let code = "let xs = [\n  1,\n  2,\n  3\n] | xs";
+        let result = response(Some(code)).unwrap();
+
+        let array_fold = result.iter().find(|r| r.kind == Some(FoldingRangeKind::Region));
+        assert!(array_fold.is_some());
+        assert_eq!(array_fold.unwrap().start_line, 0);
+        assert_eq!(array_fold.unwrap().end_line, 4);
+    }
+
+    #[test]
+    fn test_folds_consecutive_comment_lines() {
+        let code = "# First line\n# Second line\n# Third line\ndef foo(): 1;";
+        let result = response(Some(code)).unwrap();
+
+        let comment_fold = result.iter().find(|r| r.kind == Some(FoldingRangeKind::Comment));
+        assert!(comment_fold.is_some());
+        let comment_fold = comment_fold.unwrap();
+        assert_eq!(comment_fold.start_line, 0);
+        assert_eq!(comment_fold.end_line, 2);
+    }
+
+    #[test]
+    fn test_single_comment_line_not_folded() {
+        let code = "# Just one line\ndef foo(): 1;";
+        let result = response(Some(code));
+
+        let has_comment_fold = result
+            .map(|ranges| ranges.iter().any(|r| r.kind == Some(FoldingRangeKind::Comment)))
+            .unwrap_or(false);
+        assert!(!has_comment_fold);
+    }
+}

crates/mq-lsp/src/lib.rs

@@ -24,6 +24,7 @@ pub mod completions;
 pub mod document_symbol;
 pub mod error;
 pub mod execute_command;
+pub mod folding_range;
 pub mod goto_definition;
 pub mod hover;
 pub mod inlay_hints;

crates/mq-lsp/src/main.rs

@@ -10,6 +10,7 @@ pub mod completions;
 pub mod document_symbol;
 pub mod error;
 pub mod execute_command;
+pub mod folding_range;
 pub mod goto_definition;
 pub mod hover;
 pub mod inlay_hints;

crates/mq-lsp/src/server.rs

@@ -10,8 +10,8 @@ use url::Url;
 
 use crate::error::LspError;
 use crate::{
-    capabilities, code_action, completions, document_symbol, execute_command, goto_definition, hover, inlay_hints,
-    references, rename, semantic_tokens, signature_help, workspace_symbol,
+    capabilities, code_action, completions, document_symbol, execute_command, folding_range, goto_definition, hover,
+    inlay_hints, references, rename, semantic_tokens, signature_help, workspace_symbol,
 };
 use tower_lsp_server::{Client, LanguageServer, LspService, Server, jsonrpc, ls_types};
 
@@ -147,6 +147,16 @@ impl LanguageServer for Backend {
         ))
     }
 
+    async fn folding_range(
+        &self,
+        params: ls_types::FoldingRangeParams,
+    ) -> jsonrpc::Result<Option<Vec<ls_types::FoldingRange>>> {
+        let uri = params.text_document.uri;
+        let source_text = self.text_map.get(&uri.to_string()).map(|text| Arc::clone(text.value()));
+
+        Ok(folding_range::response(source_text.as_deref().map(String::as_str)))
+    }
+
     async fn symbol(
         &self,
         params: ls_types::WorkspaceSymbolParams,
@@ -850,6 +860,51 @@ mod tests {
         assert_eq!(help.active_parameter, Some(1));
     }
 
+    #[tokio::test]
+    async fn test_folding_range() {
+        let (service, _) = LspService::new(|client| Backend {
+            client,
+            hir: Arc::new(RwLock::new(mq_hir::Hir::default())),
+            source_map: RwLock::new(BiMap::new()),
+            type_env_map: DashMap::new(),
+            error_map: DashMap::new(),
+            text_map: DashMap::new(),
+            config: LspConfig::default(),
+        });
+
+        let backend = service.inner();
+        let uri = Url::parse("file:///test.mq").unwrap();
+        let code = "def foo(a):\n  let b = a + 1\n  | b;\n| foo(1)";
+
+        backend
+            .did_open(ls_types::DidOpenTextDocumentParams {
+                text_document: ls_types::TextDocumentItem {
+                    uri: to_uri(&uri),
+                    language_id: "mq".to_string(),
+                    version: 1,
+                    text: code.to_string(),
+                },
+            })
+            .await;
+
+        let result = backend
+            .folding_range(ls_types::FoldingRangeParams {
+                text_document: ls_types::TextDocumentIdentifier { uri: to_uri(&uri) },
+                work_done_progress_params: Default::default(),
+                partial_result_params: Default::default(),
+            })
+            .await;
+
+        assert!(result.is_ok());
+        let ranges = result.unwrap().unwrap();
+
+        assert!(
+            ranges
+                .iter()
+                .any(|r| r.kind == Some(ls_types::FoldingRangeKind::Region) && r.start_line == 0 && r.end_line == 2)
+        );
+    }
+
     #[tokio::test]
     async fn test_semantic_tokens() {
         let (service, _) = LspService::new(|client| Backend {
@@ -1537,6 +1592,7 @@ mod tests {
         assert!(capabilities.workspace_symbol_provider.is_some());
         assert!(capabilities.signature_help_provider.is_some());
         assert!(capabilities.diagnostic_provider.is_some());
+        assert!(capabilities.folding_range_provider.is_some());
 
         // Test shutdown
         let shutdown_result = backend.shutdown().await;