[ty] Add semantic token provider to playground (#19232)

crates/ty_wasm/src/lib.rs

@@ -338,6 +338,52 @@ impl Workspace {
             })
             .collect())
     }
+
+    #[wasm_bindgen(js_name = "semanticTokens")]
+    pub fn semantic_tokens(&self, file_id: &FileHandle) -> Result<Vec<SemanticToken>, Error> {
+        let index = line_index(&self.db, file_id.file);
+        let source = source_text(&self.db, file_id.file);
+
+        let semantic_token = ty_ide::semantic_tokens(&self.db, file_id.file, None);
+
+        let result = semantic_token
+            .iter()
+            .map(|token| SemanticToken {
+                kind: token.token_type.into(),
+                modifiers: token.modifiers.bits(),
+                range: Range::from_text_range(token.range, &index, &source, self.position_encoding),
+            })
+            .collect::<Vec<_>>();
+
+        Ok(result)
+    }
+
+    #[wasm_bindgen(js_name = "semanticTokensInRange")]
+    pub fn semantic_tokens_in_range(
+        &self,
+        file_id: &FileHandle,
+        range: Range,
+    ) -> Result<Vec<SemanticToken>, Error> {
+        let index = line_index(&self.db, file_id.file);
+        let source = source_text(&self.db, file_id.file);
+
+        let semantic_token = ty_ide::semantic_tokens(
+            &self.db,
+            file_id.file,
+            Some(range.to_text_range(&index, &source, self.position_encoding)?),
+        );
+
+        let result = semantic_token
+            .iter()
+            .map(|token| SemanticToken {
+                kind: token.token_type.into(),
+                modifiers: token.modifiers.bits(),
+                range: Range::from_text_range(token.range, &index, &source, self.position_encoding),
+            })
+            .collect::<Vec<_>>();
+
+        Ok(result)
+    }
 }
 
 pub(crate) fn into_error<E: std::fmt::Display>(err: E) -> Error {
@@ -631,6 +677,74 @@ pub struct InlayHint {
     pub position: Position,
 }
 
+#[wasm_bindgen]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SemanticToken {
+    pub kind: SemanticTokenKind,
+    pub modifiers: u32,
+    pub range: Range,
+}
+
+#[wasm_bindgen]
+impl SemanticToken {
+    pub fn kinds() -> Vec<String> {
+        ty_ide::SemanticTokenType::all()
+            .iter()
+            .map(|ty| ty.as_lsp_concept().to_string())
+            .collect()
+    }
+
+    pub fn modifiers() -> Vec<String> {
+        ty_ide::SemanticTokenModifier::all_names()
+            .iter()
+            .map(|name| (*name).to_string())
+            .collect()
+    }
+}
+
+#[wasm_bindgen]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[repr(u32)]
+pub enum SemanticTokenKind {
+    Namespace,
+    Class,
+    Parameter,
+    SelfParameter,
+    ClsParameter,
+    Variable,
+    Property,
+    Function,
+    Method,
+    Keyword,
+    String,
+    Number,
+    Decorator,
+    BuiltinConstant,
+    TypeParameter,
+}
+
+impl From<ty_ide::SemanticTokenType> for SemanticTokenKind {
+    fn from(value: ty_ide::SemanticTokenType) -> Self {
+        match value {
+            ty_ide::SemanticTokenType::Namespace => Self::Namespace,
+            ty_ide::SemanticTokenType::Class => Self::Class,
+            ty_ide::SemanticTokenType::Parameter => Self::Parameter,
+            ty_ide::SemanticTokenType::SelfParameter => Self::SelfParameter,
+            ty_ide::SemanticTokenType::ClsParameter => Self::ClsParameter,
+            ty_ide::SemanticTokenType::Variable => Self::Variable,
+            ty_ide::SemanticTokenType::Property => Self::Property,
+            ty_ide::SemanticTokenType::Function => Self::Function,
+            ty_ide::SemanticTokenType::Method => Self::Method,
+            ty_ide::SemanticTokenType::Keyword => Self::Keyword,
+            ty_ide::SemanticTokenType::String => Self::String,
+            ty_ide::SemanticTokenType::Number => Self::Number,
+            ty_ide::SemanticTokenType::Decorator => Self::Decorator,
+            ty_ide::SemanticTokenType::BuiltinConstant => Self::BuiltinConstant,
+            ty_ide::SemanticTokenType::TypeParameter => Self::TypeParameter,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 struct WasmSystem {
     fs: MemoryFileSystem,

playground/shared/src/setupMonaco.tsx

@@ -287,11 +287,7 @@ function defineAyuThemes(monaco: Monaco) {
         token: "comment",
       },
       {
-        foreground: ROCK,
+        foreground: COSMIC,
-        token: "string",
-      },
-      {
-        foreground: SUN,
         token: "keyword",
       },
       {
@@ -302,6 +298,22 @@ function defineAyuThemes(monaco: Monaco) {
         token: "tag",
         foreground: ROCK,
       },
+      {
+        foreground: ROCK,
+        token: "string",
+      },
+      {
+        token: "method",
+        foreground: SUN,
+      },
+      {
+        token: "function",
+        foreground: SUN,
+      },
+      {
+        token: "decorator",
+        foreground: SUN,
+      },
     ],
     encodedTokensColors: [],
   });
@@ -548,11 +560,11 @@ function defineAyuThemes(monaco: Monaco) {
         token: "comment",
       },
       {
-        foreground: RADIATE,
+        foreground: ELECTRON,
         token: "string",
       },
       {
-        foreground: ELECTRON,
+        foreground: CONSTELLATION,
         token: "number",
       },
       {
@@ -560,7 +572,7 @@ function defineAyuThemes(monaco: Monaco) {
         token: "identifier",
       },
       {
-        foreground: SUN,
+        foreground: RADIATE,
         token: "keyword",
       },
       {
@@ -571,6 +583,30 @@ function defineAyuThemes(monaco: Monaco) {
         foreground: ASTEROID,
         token: "delimiter",
       },
+      {
+        token: "class",
+        foreground: SUPERNOVA,
+      },
+      {
+        foreground: STARLIGHT,
+        token: "variable",
+      },
+      {
+        foreground: STARLIGHT,
+        token: "parameter",
+      },
+      {
+        token: "method",
+        foreground: SUN,
+      },
+      {
+        token: "function",
+        foreground: SUN,
+      },
+      {
+        token: "decorator",
+        foreground: SUN,
+      },
     ],
     encodedTokensColors: [],
   });

playground/ty/src/Editor/Editor.tsx

@@ -20,6 +20,7 @@ import { Theme } from "shared";
 import {
   Position as TyPosition,
   Range as TyRange,
+  SemanticToken,
   Severity,
   type Workspace,
 } from "ty_wasm";
@@ -123,6 +124,7 @@ export default function Editor({
         roundedSelection: false,
         scrollBeyondLastLine: false,
         contextmenu: true,
+        "semanticHighlighting.enabled": true,
       }}
       language={fileName.endsWith(".pyi") ? "python" : undefined}
       path={fileName}
@@ -147,7 +149,9 @@ class PlaygroundServer
     languages.HoverProvider,
     languages.InlayHintsProvider,
     languages.DocumentFormattingEditProvider,
-    languages.CompletionItemProvider
+    languages.CompletionItemProvider,
+    languages.DocumentSemanticTokensProvider,
+    languages.DocumentRangeSemanticTokensProvider
 {
   private typeDefinitionProviderDisposable: IDisposable;
   private editorOpenerDisposable: IDisposable;
@@ -155,6 +159,8 @@ class PlaygroundServer
   private inlayHintsDisposable: IDisposable;
   private formatDisposable: IDisposable;
   private completionDisposable: IDisposable;
+  private semanticTokensDisposable: IDisposable;
+  private rangeSemanticTokensDisposable: IDisposable;
 
   constructor(
     private monaco: Monaco,
@@ -174,6 +180,13 @@ class PlaygroundServer
       "python",
       this,
     );
+    this.semanticTokensDisposable =
+      monaco.languages.registerDocumentSemanticTokensProvider("python", this);
+    this.rangeSemanticTokensDisposable =
+      monaco.languages.registerDocumentRangeSemanticTokensProvider(
+        "python",
+        this,
+      );
     this.editorOpenerDisposable = monaco.editor.registerEditorOpener(this);
     this.formatDisposable =
       monaco.languages.registerDocumentFormattingEditProvider("python", this);
@@ -181,6 +194,60 @@ class PlaygroundServer
 
   triggerCharacters: string[] = ["."];
 
+  getLegend(): languages.SemanticTokensLegend {
+    return {
+      tokenTypes: SemanticToken.kinds(),
+      tokenModifiers: SemanticToken.modifiers(),
+    };
+  }
+
+  provideDocumentSemanticTokens(
+    model: editor.ITextModel,
+  ): languages.SemanticTokens | null {
+    const selectedFile = this.props.files.selected;
+
+    if (selectedFile == null) {
+      return null;
+    }
+
+    const selectedHandle = this.props.files.handles[selectedFile];
+
+    if (selectedHandle == null) {
+      return null;
+    }
+
+    const tokens = this.props.workspace.semanticTokens(selectedHandle);
+    return generateMonacoTokens(tokens, model);
+  }
+
+  releaseDocumentSemanticTokens() {}
+
+  provideDocumentRangeSemanticTokens(
+    model: editor.ITextModel,
+    range: Range,
+  ): languages.SemanticTokens | null {
+    const selectedFile = this.props.files.selected;
+
+    if (selectedFile == null) {
+      return null;
+    }
+
+    const selectedHandle = this.props.files.handles[selectedFile];
+
+    if (selectedHandle == null) {
+      return null;
+    }
+
+    const tyRange = monacoRangeToTyRange(range);
+
+    const tokens = this.props.workspace.semanticTokensInRange(
+      selectedHandle,
+      tyRange,
+    );
+
+    return generateMonacoTokens(tokens, model);
+  }
+
   provideCompletionItems(
     model: editor.ITextModel,
     position: Position,
@@ -495,6 +562,8 @@ class PlaygroundServer
     this.typeDefinitionProviderDisposable.dispose();
     this.inlayHintsDisposable.dispose();
     this.formatDisposable.dispose();
+    this.rangeSemanticTokensDisposable.dispose();
+    this.semanticTokensDisposable.dispose();
     this.completionDisposable.dispose();
   }
 }
@@ -514,3 +583,36 @@ function monacoRangeToTyRange(range: IRange): TyRange {
     new TyPosition(range.endLineNumber, range.endColumn),
   );
 }
+
+function generateMonacoTokens(
+  semantic: SemanticToken[],
+  model: editor.ITextModel,
+): languages.SemanticTokens {
+  const result = [];
+
+  let prevLine = 0;
+  let prevChar = 0;
+
+  for (const token of semantic) {
+    // Convert from 1-based to 0-based indexing for Monaco
+    const line = token.range.start.line - 1;
+    const char = token.range.start.column - 1;
+
+    const length = model.getValueLengthInRange(
+      tyRangeToMonacoRange(token.range),
+    );
+
+    result.push(
+      line - prevLine,
+      prevLine === line ? char - prevChar : char,
+      length,
+      token.kind,
+      token.modifiers,
+    );
+
+    prevLine = line;
+    prevChar = char;
+  }
+
+  return { data: Uint32Array.from(result) };
+}