[ty] Added support for document highlights in playground. (#19540)

This PR adds support for the "document highlights" feature in the ty
playground.

---------

Co-authored-by: UnboundVariable <unbound@gmail.com>

crates/ty_ide/src/goto.rs

@@ -308,7 +308,6 @@ impl GotoTarget<'_> {
                 }
             }
 
-            // TODO: Handle string literals that map to TypedDict fields
             _ => None,
         }
     }

crates/ty_wasm/src/lib.rs

@@ -15,8 +15,8 @@ use ruff_python_formatter::formatted_file;
 use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
 use ruff_text_size::{Ranged, TextSize};
 use ty_ide::{
-    MarkupKind, NavigationTargets, RangedValue, goto_declaration, goto_definition, goto_references,
-    goto_type_definition, hover, inlay_hints, signature_help,
+    MarkupKind, NavigationTargets, RangedValue, document_highlights, goto_declaration,
+    goto_definition, goto_references, goto_type_definition, hover, inlay_hints, signature_help,
 };
 use ty_project::metadata::options::Options;
 use ty_project::metadata::value::ValueSource;
@@ -528,6 +528,34 @@ impl Workspace {
                 .and_then(|s| u32::try_from(s).ok()),
         }))
     }
+
+    #[wasm_bindgen(js_name = "documentHighlights")]
+    pub fn document_highlights(
+        &self,
+        file_id: &FileHandle,
+        position: Position,
+    ) -> Result<Vec<DocumentHighlight>, Error> {
+        let source = source_text(&self.db, file_id.file);
+        let index = line_index(&self.db, file_id.file);
+
+        let offset = position.to_text_size(&source, &index, self.position_encoding)?;
+
+        let Some(targets) = document_highlights(&self.db, file_id.file, offset) else {
+            return Ok(Vec::new());
+        };
+
+        Ok(targets
+            .into_iter()
+            .map(|target| DocumentHighlight {
+                range: Range::from_file_range(
+                    &self.db,
+                    target.file_range(),
+                    self.position_encoding,
+                ),
+                kind: target.kind().into(),
+            })
+            .collect())
+    }
 }
 
 pub(crate) fn into_error<E: std::fmt::Display>(err: E) -> Error {
@@ -954,6 +982,33 @@ pub struct ParameterInformation {
     pub documentation: Option<String>,
 }
 
+#[wasm_bindgen]
+pub struct DocumentHighlight {
+    #[wasm_bindgen(readonly)]
+    pub range: Range,
+
+    #[wasm_bindgen(readonly)]
+    pub kind: DocumentHighlightKind,
+}
+
+#[wasm_bindgen]
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum DocumentHighlightKind {
+    Text = 1,
+    Read = 2,
+    Write = 3,
+}
+
+impl From<ty_ide::ReferenceKind> for DocumentHighlightKind {
+    fn from(kind: ty_ide::ReferenceKind) -> Self {
+        match kind {
+            ty_ide::ReferenceKind::Read => DocumentHighlightKind::Read,
+            ty_ide::ReferenceKind::Write => DocumentHighlightKind::Write,
+            ty_ide::ReferenceKind::Other => DocumentHighlightKind::Text,
+        }
+    }
+}
+
 #[wasm_bindgen]
 impl SemanticToken {
     pub fn kinds() -> Vec<String> {

playground/ty/src/Editor/Editor.tsx

@@ -24,6 +24,8 @@ import {
   Severity,
   type Workspace,
   CompletionKind,
+  DocumentHighlight,
+  DocumentHighlightKind,
 } from "ty_wasm";
 import { FileId, ReadonlyFiles } from "../Playground";
 import { isPythonFile } from "./Files";
@@ -156,7 +158,8 @@ class PlaygroundServer
     languages.CompletionItemProvider,
     languages.DocumentSemanticTokensProvider,
     languages.DocumentRangeSemanticTokensProvider,
-    languages.SignatureHelpProvider
+    languages.SignatureHelpProvider,
+    languages.DocumentHighlightProvider
 {
   private typeDefinitionProviderDisposable: IDisposable;
   private declarationProviderDisposable: IDisposable;
@@ -170,6 +173,7 @@ class PlaygroundServer
   private semanticTokensDisposable: IDisposable;
   private rangeSemanticTokensDisposable: IDisposable;
   private signatureHelpDisposable: IDisposable;
+  private documentHighlightDisposable: IDisposable;
 
   constructor(
     private monaco: Monaco,
@@ -207,6 +211,8 @@ class PlaygroundServer
       monaco.languages.registerDocumentFormattingEditProvider("python", this);
     this.signatureHelpDisposable =
       monaco.languages.registerSignatureHelpProvider("python", this);
+    this.documentHighlightDisposable =
+      monaco.languages.registerDocumentHighlightProvider("python", this);
   }
 
   triggerCharacters: string[] = ["."];
@@ -365,6 +371,36 @@ class PlaygroundServer
     };
   }
 
+  provideDocumentHighlights(
+    model: editor.ITextModel,
+    position: Position,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _token: CancellationToken,
+  ): languages.ProviderResult<languages.DocumentHighlight[]> {
+    const workspace = this.props.workspace;
+    const selectedFile = this.props.files.selected;
+
+    if (selectedFile == null) {
+      return;
+    }
+
+    const selectedHandle = this.props.files.handles[selectedFile];
+
+    if (selectedHandle == null) {
+      return;
+    }
+
+    const highlights = workspace.documentHighlights(
+      selectedHandle,
+      new TyPosition(position.lineNumber, position.column),
+    );
+
+    return highlights.map((highlight: DocumentHighlight) => ({
+      range: tyRangeToMonacoRange(highlight.range),
+      kind: mapDocumentHighlightKind(highlight.kind),
+    }));
+  }
+
   provideInlayHints(
     _model: editor.ITextModel,
     range: Range,
@@ -707,6 +743,7 @@ class PlaygroundServer
     this.semanticTokensDisposable.dispose();
     this.completionDisposable.dispose();
     this.signatureHelpDisposable.dispose();
+    this.documentHighlightDisposable.dispose();
   }
 }
 
@@ -836,3 +873,18 @@ function mapCompletionKind(kind: CompletionKind): CompletionItemKind {
       return CompletionItemKind.TypeParameter;
   }
 }
+
+function mapDocumentHighlightKind(
+  kind: DocumentHighlightKind,
+): languages.DocumentHighlightKind {
+  switch (kind) {
+    case DocumentHighlightKind.Text:
+      return languages.DocumentHighlightKind.Text;
+    case DocumentHighlightKind.Read:
+      return languages.DocumentHighlightKind.Read;
+    case DocumentHighlightKind.Write:
+      return languages.DocumentHighlightKind.Write;
+    default:
+      return languages.DocumentHighlightKind.Text;
+  }
+}