feat(lsp): implement textDocument/foldingRange (#9900)

Co-authored-by: Kitson Kelly <me@kitsonkelly.com>

cli/lsp/capabilities.rs

@@ -11,6 +11,7 @@ use lspower::lsp::CodeActionOptions;
 use lspower::lsp::CodeActionProviderCapability;
 use lspower::lsp::CodeLensOptions;
 use lspower::lsp::CompletionOptions;
+use lspower::lsp::FoldingRangeProviderCapability;
 use lspower::lsp::HoverProviderCapability;
 use lspower::lsp::ImplementationProviderCapability;
 use lspower::lsp::OneOf;
@@ -108,7 +109,7 @@ pub fn server_capabilities(
     selection_range_provider: Some(SelectionRangeProviderCapability::Simple(
       true,
     )),
-    folding_range_provider: None,
+    folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
     rename_provider: Some(OneOf::Left(true)),
     document_link_provider: None,
     color_provider: None,

cli/lsp/config.rs

@@ -13,6 +13,7 @@ pub struct ClientCapabilities {
   pub status_notification: bool,
   pub workspace_configuration: bool,
   pub workspace_did_change_watched_files: bool,
+  pub line_folding_only: bool,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -125,5 +126,13 @@ impl Config {
         .and_then(|it| it.dynamic_registration)
         .unwrap_or(false);
     }
+
+    if let Some(text_document) = &capabilities.text_document {
+      self.client_capabilities.line_folding_only = text_document
+        .folding_range
+        .as_ref()
+        .and_then(|it| it.line_folding_only)
+        .unwrap_or(false);
+    }
   }
 }

cli/lsp/language_server.rs

@@ -228,6 +228,25 @@ impl Inner {
     maybe_line_index
   }
 
+  // TODO(@kitsonk) we really should find a better way to just return the
+  // content as a `&str`, or be able to get the byte at a particular offset
+  // which is all that this API that is consuming it is trying to do at the
+  // moment
+  /// Searches already cached assets and documents and returns its text
+  /// content. If not found, `None` is returned.
+  fn get_text_content(&self, specifier: &ModuleSpecifier) -> Option<String> {
+    if specifier.scheme() == "asset" {
+      self
+        .assets
+        .get(specifier)
+        .map(|o| o.clone().map(|a| a.text))?
+    } else if self.documents.contains_key(specifier) {
+      self.documents.content(specifier).unwrap()
+    } else {
+      self.sources.get_source(specifier)
+    }
+  }
+
   async fn get_navigation_tree(
     &mut self,
     specifier: &ModuleSpecifier,
@@ -1515,6 +1534,63 @@ impl Inner {
     Ok(result)
   }
 
+  async fn folding_range(
+    &self,
+    params: FoldingRangeParams,
+  ) -> LspResult<Option<Vec<FoldingRange>>> {
+    if !self.enabled() {
+      return Ok(None);
+    }
+    let mark = self.performance.mark("folding_range");
+    let specifier = self.url_map.normalize_url(&params.text_document.uri);
+
+    let line_index =
+      if let Some(line_index) = self.get_line_index_sync(&specifier) {
+        line_index
+      } else {
+        return Err(LspError::invalid_params(format!(
+          "An unexpected specifier ({}) was provided.",
+          specifier
+        )));
+      };
+
+    let req = tsc::RequestMethod::GetOutliningSpans(specifier.clone());
+    let outlining_spans: Vec<tsc::OutliningSpan> = self
+      .ts_server
+      .request(self.snapshot(), req)
+      .await
+      .map_err(|err| {
+        error!("Failed to request to tsserver {}", err);
+        LspError::invalid_request()
+      })?;
+
+    let response = if !outlining_spans.is_empty() {
+      let text_content =
+        self.get_text_content(&specifier).ok_or_else(|| {
+          LspError::invalid_params(format!(
+            "An unexpected specifier ({}) was provided.",
+            specifier
+          ))
+        })?;
+      Some(
+        outlining_spans
+          .iter()
+          .map(|span| {
+            span.to_folding_range(
+              &line_index,
+              text_content.as_str().as_bytes(),
+              self.config.client_capabilities.line_folding_only,
+            )
+          })
+          .collect::<Vec<FoldingRange>>(),
+      )
+    } else {
+      None
+    };
+    self.performance.measure(mark);
+    Ok(response)
+  }
+
   async fn rename(
     &mut self,
     params: RenameParams,
@@ -1840,6 +1916,13 @@ impl lspower::LanguageServer for LanguageServer {
     self.0.lock().await.goto_implementation(params).await
   }
 
+  async fn folding_range(
+    &self,
+    params: FoldingRangeParams,
+  ) -> LspResult<Option<Vec<FoldingRange>>> {
+    self.0.lock().await.folding_range(params).await
+  }
+
   async fn rename(
     &self,
     params: RenameParams,
@@ -2425,6 +2508,54 @@ mod tests {
     );
   }
 
+  #[tokio::test]
+  async fn test_folding_range() {
+    let mut harness = LspTestHarness::new(vec![
+      ("initialize_request.json", LspResponse::RequestAny),
+      ("initialized_notification.json", LspResponse::None),
+      (
+        "folding_range_did_open_notification.json",
+        LspResponse::None,
+      ),
+      (
+        "folding_range_request.json",
+        LspResponse::Request(
+          2,
+          json!([
+            {
+              "startLine": 0,
+              "endLine": 12,
+              "kind": "region"
+            },
+            {
+              "startLine": 1,
+              "endLine": 3,
+              "kind": "comment"
+            },
+            {
+              "startLine": 4,
+              "endLine": 10
+            },
+            {
+              "startLine": 5,
+              "endLine": 9
+            },
+            {
+              "startLine": 6,
+              "endLine": 7
+            }
+          ]),
+        ),
+      ),
+      (
+        "shutdown_request.json",
+        LspResponse::Request(3, json!(null)),
+      ),
+      ("exit_notification.json", LspResponse::None),
+    ]);
+    harness.run().await;
+  }
+
   #[tokio::test]
   async fn test_rename() {
     let mut harness = LspTestHarness::new(vec![

cli/lsp/tsc.rs

@@ -35,10 +35,10 @@ use log::warn;
 use lspower::lsp;
 use regex::Captures;
 use regex::Regex;
-use std::borrow::Cow;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::thread;
+use std::{borrow::Cow, cmp};
 use text_size::TextSize;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
@@ -1227,6 +1227,91 @@ impl CompletionEntry {
   }
 }
 
+#[derive(Debug, Deserialize)]
+pub enum OutliningSpanKind {
+  #[serde(rename = "comment")]
+  Comment,
+  #[serde(rename = "region")]
+  Region,
+  #[serde(rename = "code")]
+  Code,
+  #[serde(rename = "imports")]
+  Imports,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct OutliningSpan {
+  text_span: TextSpan,
+  hint_span: TextSpan,
+  banner_text: String,
+  auto_collapse: bool,
+  kind: OutliningSpanKind,
+}
+
+const FOLD_END_PAIR_CHARACTERS: &[u8] = &[b'}', b']', b')', b'`'];
+
+impl OutliningSpan {
+  pub fn to_folding_range(
+    &self,
+    line_index: &LineIndex,
+    content: &[u8],
+    line_folding_only: bool,
+  ) -> lsp::FoldingRange {
+    let range = self.text_span.to_range(line_index);
+    lsp::FoldingRange {
+      start_line: range.start.line,
+      start_character: if line_folding_only {
+        None
+      } else {
+        Some(range.start.character)
+      },
+      end_line: self.adjust_folding_end_line(
+        &range,
+        line_index,
+        content,
+        line_folding_only,
+      ),
+      end_character: if line_folding_only {
+        None
+      } else {
+        Some(range.end.character)
+      },
+      kind: self.get_folding_range_kind(&self.kind),
+    }
+  }
+
+  fn adjust_folding_end_line(
+    &self,
+    range: &lsp::Range,
+    line_index: &LineIndex,
+    content: &[u8],
+    line_folding_only: bool,
+  ) -> u32 {
+    if line_folding_only && range.end.character > 0 {
+      let offset_end: usize = line_index.offset(range.end).unwrap().into();
+      let fold_end_char = content[offset_end - 1];
+      if FOLD_END_PAIR_CHARACTERS.contains(&fold_end_char) {
+        return cmp::max(range.end.line - 1, range.start.line);
+      }
+    }
+
+    range.end.line
+  }
+
+  fn get_folding_range_kind(
+    &self,
+    span_kind: &OutliningSpanKind,
+  ) -> Option<lsp::FoldingRangeKind> {
+    match span_kind {
+      OutliningSpanKind::Comment => Some(lsp::FoldingRangeKind::Comment),
+      OutliningSpanKind::Region => Some(lsp::FoldingRangeKind::Region),
+      OutliningSpanKind::Imports => Some(lsp::FoldingRangeKind::Imports),
+      _ => None,
+    }
+  }
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct SignatureHelpItems {
@@ -1859,6 +1944,8 @@ pub enum RequestMethod {
   GetImplementation((ModuleSpecifier, u32)),
   /// Get a "navigation tree" for a specifier.
   GetNavigationTree(ModuleSpecifier),
+  /// Get outlining spans for a specifier.
+  GetOutliningSpans(ModuleSpecifier),
   /// Return quick info at position (hover information).
   GetQuickInfo((ModuleSpecifier, u32)),
   /// Get document references for a specific position.
@@ -1967,6 +2054,11 @@ impl RequestMethod {
         "method": "getNavigationTree",
         "specifier": specifier,
       }),
+      RequestMethod::GetOutliningSpans(specifier) => json!({
+        "id": id,
+        "method": "getOutliningSpans",
+        "specifier": specifier,
+      }),
       RequestMethod::GetQuickInfo((specifier, position)) => json!({
         "id": id,
         "method": "getQuickInfo",

cli/tests/lsp/folding_range_did_open_notification.json

@@ -0,0 +1,12 @@
+{
+  "jsonrpc": "2.0",
+  "method": "textDocument/didOpen",
+  "params": {
+    "textDocument": {
+      "uri": "file:///a/file.ts",
+      "languageId": "typescript",
+      "version": 1,
+      "text": "// #region 1\n/*\n * Some comment\n */\nclass Foo {\n  bar(a, b) {\n    if (a === b) {\n      return true;\n    }\n    return false;\n  }\n}\n// #endregion"
+    }
+  }
+}

cli/tests/lsp/folding_range_request.json

@@ -0,0 +1,10 @@
+{
+  "jsonrpc": "2.0",
+  "id": 2,
+  "method": "textDocument/foldingRange",
+  "params": {
+    "textDocument": {
+      "uri": "file:///a/file.ts"
+    }
+  }
+}

cli/tests/lsp/initialize_request.json

@@ -37,6 +37,9 @@
             ]
           }
         },
+        "foldingRange": {
+          "lineFoldingOnly": true
+        },
         "synchronization": {
           "dynamicRegistration": true,
           "willSave": true,

cli/tsc/99_main_compiler.js

@@ -675,6 +675,14 @@ delete Object.prototype.__proto__;
           languageService.getNavigationTree(request.specifier),
         );
       }
+      case "getOutliningSpans": {
+        return respond(
+          id,
+          languageService.getOutliningSpans(
+            request.specifier,
+          ),
+        );
+      }
       case "getQuickInfo": {
         return respond(
           id,

cli/tsc/compiler.d.ts

@@ -58,6 +58,7 @@ declare global {
     | GetDocumentHighlightsRequest
     | GetImplementationRequest
     | GetNavigationTree
+    | GetOutliningSpans
     | GetQuickInfoRequest
     | GetReferencesRequest
     | GetSignatureHelpItemsRequest
@@ -151,6 +152,11 @@ declare global {
     specifier: string;
   }
 
+  interface GetOutliningSpans extends BaseLanguageServerRequest {
+    method: "getOutliningSpans";
+    specifier: string;
+  }
+
   interface GetQuickInfoRequest extends BaseLanguageServerRequest {
     method: "getQuickInfo";
     specifier: string;