feat(lsp): Implement textDocument/semanticTokens/full (#10233)

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

cli/lsp/capabilities.rs

@@ -18,13 +18,18 @@ use lspower::lsp::ImplementationProviderCapability;
 use lspower::lsp::OneOf;
 use lspower::lsp::SaveOptions;
 use lspower::lsp::SelectionRangeProviderCapability;
+use lspower::lsp::SemanticTokensFullOptions;
+use lspower::lsp::SemanticTokensOptions;
+use lspower::lsp::SemanticTokensServerCapabilities;
 use lspower::lsp::ServerCapabilities;
 use lspower::lsp::SignatureHelpOptions;
 use lspower::lsp::TextDocumentSyncCapability;
 use lspower::lsp::TextDocumentSyncKind;
 use lspower::lsp::TextDocumentSyncOptions;
 use lspower::lsp::WorkDoneProgressOptions;
 
+use super::semantic_tokens::get_legend;
+
 fn code_action_capabilities(
   client_capabilities: &ClientCapabilities,
 ) -> CodeActionProviderCapability {
@@ -116,7 +121,16 @@ pub fn server_capabilities(
     color_provider: None,
     execute_command_provider: None,
     call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
-    semantic_tokens_provider: None,
+    semantic_tokens_provider: Some(
+      SemanticTokensServerCapabilities::SemanticTokensOptions(
+        SemanticTokensOptions {
+          legend: get_legend(),
+          range: Some(true),
+          full: Some(SemanticTokensFullOptions::Bool(true)),
+          ..Default::default()
+        },
+      ),
+    ),
     workspace: None,
     experimental: None,
     linked_editing_range_provider: None,

cli/lsp/language_server.rs

@@ -1961,6 +1961,99 @@ impl Inner {
     Ok(Some(selection_ranges))
   }
 
+  async fn semantic_tokens_full(
+    &self,
+    params: SemanticTokensParams,
+  ) -> LspResult<Option<SemanticTokensResult>> {
+    if !self.enabled() {
+      return Ok(None);
+    }
+    let mark = self.performance.mark("semantic_tokens_full");
+    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::GetEncodedSemanticClassifications((
+      specifier.clone(),
+      tsc::TextSpan {
+        start: 0,
+        length: line_index.text_content_length_utf16().into(),
+      },
+    ));
+    let semantic_classification: tsc::Classifications = self
+      .ts_server
+      .request(self.snapshot(), req)
+      .await
+      .map_err(|err| {
+        error!("Failed to request to tsserver {}", err);
+        LspError::invalid_request()
+      })?;
+
+    let semantic_tokens: SemanticTokens =
+      semantic_classification.to_semantic_tokens(&line_index);
+    let response = if !semantic_tokens.data.is_empty() {
+      Some(SemanticTokensResult::Tokens(semantic_tokens))
+    } else {
+      None
+    };
+    self.performance.measure(mark);
+    Ok(response)
+  }
+
+  async fn semantic_tokens_range(
+    &self,
+    params: SemanticTokensRangeParams,
+  ) -> LspResult<Option<SemanticTokensRangeResult>> {
+    if !self.enabled() {
+      return Ok(None);
+    }
+    let mark = self.performance.mark("semantic_tokens_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 start = line_index.offset_tsc(params.range.start)?;
+    let length = line_index.offset_tsc(params.range.end)? - start;
+    let req = tsc::RequestMethod::GetEncodedSemanticClassifications((
+      specifier.clone(),
+      tsc::TextSpan { start, length },
+    ));
+    let semantic_classification: tsc::Classifications = self
+      .ts_server
+      .request(self.snapshot(), req)
+      .await
+      .map_err(|err| {
+        error!("Failed to request to tsserver {}", err);
+        LspError::invalid_request()
+      })?;
+
+    let semantic_tokens: SemanticTokens =
+      semantic_classification.to_semantic_tokens(&line_index);
+    let response = if !semantic_tokens.data.is_empty() {
+      Some(SemanticTokensRangeResult::Tokens(semantic_tokens))
+    } else {
+      None
+    };
+    self.performance.measure(mark);
+    Ok(response)
+  }
+
   async fn signature_help(
     &self,
     params: SignatureHelpParams,
@@ -2200,6 +2293,20 @@ impl lspower::LanguageServer for LanguageServer {
     self.0.lock().await.selection_range(params).await
   }
 
+  async fn semantic_tokens_full(
+    &self,
+    params: SemanticTokensParams,
+  ) -> LspResult<Option<SemanticTokensResult>> {
+    self.0.lock().await.semantic_tokens_full(params).await
+  }
+
+  async fn semantic_tokens_range(
+    &self,
+    params: SemanticTokensRangeParams,
+  ) -> LspResult<Option<SemanticTokensRangeResult>> {
+    self.0.lock().await.semantic_tokens_range(params).await
+  }
+
   async fn signature_help(
     &self,
     params: SignatureHelpParams,
@@ -3539,6 +3646,43 @@ mod tests {
     harness.run().await;
   }
 
+  #[tokio::test]
+  #[rustfmt::skip]
+  async fn test_semantic_tokens() {
+    let mut harness = LspTestHarness::new(vec![
+      (LspFixture::Path("initialize_request.json"), LspResponse::RequestAny),
+      (LspFixture::Path("initialized_notification.json"), LspResponse::None),
+      (
+        LspFixture::Path("semantic_tokens_did_open_notification.json"),
+        LspResponse::None,
+      ),
+      (
+        LspFixture::Path("semantic_tokens_full_request.json"),
+        LspResponse::Request(
+          2,
+          json!({
+            "data": [0, 5, 6, 1, 1, 0, 9, 6, 8, 9, 0, 8, 6, 8, 9, 2, 15 ,3, 10 ,5, 0, 4, 1, 6, 1, 0, 12 ,7, 2, 16 ,1, 8, 1, 7, 41 ,0, 4, 1, 6, 0, 0, 2, 5, 11 ,16 ,1, 9, 1, 7, 40 ,3, 10 ,4, 2, 1, 1, 11 ,1, 9, 9, 1, 2, 3, 11 ,1, 3, 6, 3, 0, 1, 0, 15 ,4, 2, 0, 1, 30 ,1, 6, 9, 1, 2, 3, 11 ,1, 1, 9, 9, 9, 3, 0, 16 ,3, 0, 0, 1, 17 ,12 ,11 ,3, 0, 24 ,3, 0, 0, 0, 4, 9, 9, 2]
+          }),
+        ),
+      ),
+      (
+        LspFixture::Path("semantic_tokens_range_request.json"),
+        LspResponse::Request(
+          4,
+          json!({
+            "data": [0, 5, 6, 1, 1, 0, 9, 6, 8, 9, 0, 8, 6, 8, 9, 2, 15 ,3, 10 ,5, 0, 4, 1, 6, 1, 0, 12 ,7, 2, 16 ,1, 8, 1, 7, 41 ,0, 4, 1, 6, 0, 0, 2, 5, 11 ,16 ,1, 9, 1, 7, 40]
+          }),
+        ),
+      ),
+      (
+        LspFixture::Path("shutdown_request.json"),
+        LspResponse::Request(3, json!(null)),
+      ),
+      (LspFixture::Path("exit_notification.json"), LspResponse::None),
+    ]);
+    harness.run().await;
+  }
+
   #[tokio::test]
   async fn test_code_lens_request() {
     let mut harness = LspTestHarness::new(vec![

cli/lsp/mod.rs

@@ -13,6 +13,7 @@ pub(crate) mod language_server;
 mod path_to_regex;
 mod performance;
 mod registries;
+mod semantic_tokens;
 mod sources;
 mod text;
 mod tsc;