feat(LSP): folding ranges

tooling/lsp/src/lib.rs

@@ -14,8 +14,9 @@ use std::{
 
 use acvm::{BlackBoxFunctionSolver, FieldElement};
 use async_lsp::lsp_types::request::{
-    CodeActionRequest, Completion, DocumentSymbolRequest, HoverRequest, InlayHintRequest,
-    PrepareRenameRequest, References, Rename, SignatureHelpRequest, WorkspaceSymbolRequest,
+    CodeActionRequest, Completion, DocumentSymbolRequest, FoldingRangeRequest, HoverRequest,
+    InlayHintRequest, PrepareRenameRequest, References, Rename, SignatureHelpRequest,
+    WorkspaceSymbolRequest,
 };
 use async_lsp::{
     AnyEvent, AnyNotification, AnyRequest, ClientSocket, Error, LspService, ResponseError,
@@ -79,7 +80,7 @@ use types::{NargoTest, NargoTestId, Position, Range, Url, notification, request}
 use with_file::parsed_module_with_file;
 
 use crate::{
-    requests::{on_expand_request, on_std_source_code_request},
+    requests::{on_expand_request, on_folding_range_request, on_std_source_code_request},
     types::request::{NargoExpand, NargoStdSourceCode},
 };
 
@@ -176,6 +177,7 @@ impl NargoLspService {
             .request::<SignatureHelpRequest, _>(on_signature_help_request)
             .request::<CodeActionRequest, _>(on_code_action_request)
             .request::<WorkspaceSymbolRequest, _>(on_workspace_symbol_request)
+            .request::<FoldingRangeRequest, _>(on_folding_range_request)
             .request::<NargoExpand, _>(on_expand_request)
             .request::<NargoStdSourceCode, _>(on_std_source_code_request)
             .notification::<notification::Initialized>(on_initialized)

tooling/lsp/src/requests/folding_range.rs

@@ -0,0 +1,47 @@
+use std::future;
+
+use async_lsp::{
+    ResponseError,
+    lsp_types::{FoldingRange, FoldingRangeParams, Position, TextDocumentPositionParams},
+};
+
+use crate::{
+    LspState,
+    requests::{
+        folding_range::{comments_collector::CommentsCollector, nodes_collector::NodesCollector},
+        process_request,
+    },
+};
+
+mod comments_collector;
+mod nodes_collector;
+#[cfg(test)]
+mod tests;
+
+pub(crate) fn on_folding_range_request(
+    state: &mut LspState,
+    params: FoldingRangeParams,
+) -> impl Future<Output = Result<Option<Vec<FoldingRange>>, ResponseError>> + use<> {
+    let text_document_position_params = TextDocumentPositionParams {
+        text_document: params.text_document.clone(),
+        position: Position { line: 0, character: 0 },
+    };
+
+    let result = process_request(state, text_document_position_params, |args| {
+        let file_id = args.location.file;
+        let file = args.files.get_file(file_id).unwrap();
+        let source = file.source();
+
+        let comments_collector = CommentsCollector::new(file_id, args.files);
+        let mut ranges = comments_collector.collect(source);
+
+        let nodes_collector = NodesCollector::new(file_id, args.files);
+        let node_ranges = nodes_collector.collect(source);
+
+        ranges.extend(node_ranges);
+
+        Some(ranges)
+    });
+
+    future::ready(result)
+}

tooling/lsp/src/requests/folding_range/comments_collector.rs

@@ -0,0 +1,113 @@
+use async_lsp::lsp_types::{self, FoldingRange, FoldingRangeKind};
+use fm::{FileId, FileMap};
+use noirc_frontend::{
+    lexer::Lexer,
+    token::{DocStyle, Token},
+};
+
+use crate::requests::to_lsp_location;
+
+pub(super) struct CommentsCollector<'files> {
+    file_id: FileId,
+    files: &'files FileMap,
+    ranges: Vec<FoldingRange>,
+    current_line_comment_group: Option<LineCommentGroup>,
+}
+
+struct LineCommentGroup {
+    start: lsp_types::Location,
+    end: lsp_types::Location,
+    doc_style: Option<DocStyle>,
+}
+
+impl<'files> CommentsCollector<'files> {
+    pub(super) fn new(file_id: FileId, files: &'files FileMap) -> Self {
+        Self { file_id, files, ranges: Vec::new(), current_line_comment_group: None }
+    }
+
+    pub(super) fn collect(mut self, source: &str) -> Vec<FoldingRange> {
+        let lexer = Lexer::new(source, self.file_id).skip_comments(false);
+
+        for token in lexer {
+            let Ok(token) = token else {
+                continue;
+            };
+
+            let location = token.location();
+
+            match token.into_token() {
+                Token::BlockComment(..) => {
+                    self.push_current_line_comment_group();
+
+                    let Some(location) = to_lsp_location(self.files, self.file_id, location.span)
+                    else {
+                        continue;
+                    };
+
+                    // Block comments are never grouped with other comments
+                    self.ranges.push(FoldingRange {
+                        start_line: location.range.start.line,
+                        start_character: None,
+                        end_line: location.range.end.line,
+                        end_character: None,
+                        kind: Some(FoldingRangeKind::Comment),
+                        collapsed_text: None,
+                    });
+                }
+                Token::LineComment(_, doc_style) => {
+                    let Some(location) = to_lsp_location(self.files, self.file_id, location.span)
+                    else {
+                        continue;
+                    };
+
+                    if let Some(group) = &mut self.current_line_comment_group {
+                        // Keep grouping while the line comment style is the same and they are consecutive lines
+                        if group.doc_style == doc_style
+                            && location.range.start.line - group.end.range.end.line <= 1
+                        {
+                            group.end = location;
+                            continue;
+                        }
+
+                        // A new group starts, so push the current one
+                        Self::push_line_comment_group(group, &mut self.ranges);
+                    }
+
+                    let start = location.clone();
+                    self.current_line_comment_group =
+                        Some(LineCommentGroup { start, end: location, doc_style });
+                }
+                _ => {
+                    self.push_current_line_comment_group();
+                }
+            }
+        }
+
+        self.push_current_line_comment_group();
+
+        self.ranges
+    }
+
+    fn push_current_line_comment_group(&mut self) {
+        if let Some(group) = self.current_line_comment_group.take() {
+            Self::push_line_comment_group(&group, &mut self.ranges);
+        }
+    }
+
+    fn push_line_comment_group(group: &LineCommentGroup, ranges: &mut Vec<FoldingRange>) {
+        let start_line = group.start.range.start.line;
+        let end_line = group.end.range.end.line;
+        if start_line == end_line {
+            return;
+        }
+
+        ranges.push(FoldingRange {
+            start_line,
+            start_character: None,
+            end_line,
+            end_character: None,
+            kind: Some(FoldingRangeKind::Comment),
+            collapsed_text: None,
+        });
+    }
+}

tooling/lsp/src/requests/folding_range/nodes_collector.rs

@@ -0,0 +1,123 @@
+use async_lsp::lsp_types::{self, FoldingRange, FoldingRangeKind};
+use fm::{FileId, FileMap};
+use noirc_errors::Span;
+use noirc_frontend::ast::{ItemVisibility, ModuleDeclaration, UseTree, Visitor};
+
+use crate::requests::to_lsp_location;
+
+pub(super) struct NodesCollector<'files> {
+    file_id: FileId,
+    files: &'files FileMap,
+    ranges: Vec<FoldingRange>,
+    module_group: Option<Group>,
+    use_group: Option<Group>,
+}
+
+struct Group {
+    start: lsp_types::Location,
+    end: lsp_types::Location,
+    count: usize,
+}
+
+impl<'files> NodesCollector<'files> {
+    pub(super) fn new(file_id: FileId, files: &'files FileMap) -> Self {
+        Self { file_id, files, ranges: Vec::new(), module_group: None, use_group: None }
+    }
+
+    pub(super) fn collect(mut self, source: &str) -> Vec<FoldingRange> {
+        let (parsed_module, _errors) = noirc_frontend::parse_program(source, self.file_id);
+        parsed_module.accept(&mut self);
+
+        if let Some(group) = &self.module_group {
+            Self::push_group(group, None, &mut self.ranges);
+        }
+
+        if let Some(group) = &self.use_group {
+            Self::push_group(group, Some(FoldingRangeKind::Imports), &mut self.ranges);
+        }
+
+        self.ranges
+    }
+
+    fn push_group(group: &Group, kind: Option<FoldingRangeKind>, ranges: &mut Vec<FoldingRange>) {
+        if group.count == 1 {
+            return;
+        }
+
+        let start_line = group.start.range.start.line;
+        let end_line = group.end.range.end.line;
+        if start_line == end_line {
+            return;
+        }
+
+        ranges.push(FoldingRange {
+            start_line,
+            start_character: None,
+            end_line,
+            end_character: None,
+            kind,
+            collapsed_text: None,
+        });
+    }
+}
+
+impl Visitor for NodesCollector<'_> {
+    fn visit_module_declaration(&mut self, _: &ModuleDeclaration, span: Span) {
+        let Some(location) = to_lsp_location(self.files, self.file_id, span) else {
+            return;
+        };
+
+        if let Some(group) = &mut self.module_group {
+            if location.range.start.line - group.end.range.end.line <= 1 {
+                group.end = location;
+                group.count += 1;
+                return;
+            }
+
+            Self::push_group(group, None, &mut self.ranges);
+        }
+
+        self.module_group = Some(Group { start: location.clone(), end: location, count: 1 });
+    }
+
+    fn visit_import(&mut self, _: &UseTree, span: Span, _visibility: ItemVisibility) -> bool {
+        let Some(location) = to_lsp_location(self.files, self.file_id, span) else {
+            return true;
+        };
+
+        if let Some(group) = &mut self.use_group {
+            if location.range.start.line - group.end.range.end.line <= 1 {
+                group.end = location;
+                group.count += 1;
+                return true;
+            }
+
+            Self::push_group(group, Some(FoldingRangeKind::Imports), &mut self.ranges);
+        }
+
+        self.use_group = Some(Group { start: location.clone(), end: location, count: 1 });
+
+        true
+    }
+
+    fn visit_use_tree(&mut self, use_tree: &UseTree) -> bool {
+        let Some(location) = to_lsp_location(self.files, self.file_id, use_tree.location.span)
+        else {
+            return true;
+        };
+
+        // If the use tree spans multiple lines, we can fold it
+        if location.range.start.line != location.range.end.line {
+            self.ranges.push(FoldingRange {
+                start_line: location.range.start.line,
+                start_character: None,
+                end_line: location.range.end.line,
+                end_character: None,
+                kind: Some(FoldingRangeKind::Imports),
+                collapsed_text: None,
+            });
+        }
+
+        true
+    }
+}