[red-knot] Add 'Goto type definition' to the playground (#17055)

## Summary

This PR adds Goto type definition to the playground, using the same
infrastructure as the LSP.


The main *challenge* with implementing this feature was that the editor
can now participate in which tab is open.

## Known limitations

The same as for the LSP. Most notably, navigating to types defined in
typeshed isn't supported.

## Test Plan


https://github.com/user-attachments/assets/22dad7c8-7ac7-463f-b066-5d5b2c45d1fe

Cargo.lock

@@ -2665,6 +2665,7 @@ dependencies = [
  "getrandom 0.3.2",
  "js-sys",
  "log",
+ "red_knot_ide",
  "red_knot_project",
  "red_knot_python_semantic",
  "ruff_db",

crates/red_knot_ide/src/lib.rs

@@ -21,6 +21,12 @@ pub struct RangedValue<T> {
     pub value: T,
 }
 
+impl<T> RangedValue<T> {
+    pub fn file_range(&self) -> FileRange {
+        self.range
+    }
+}
+
 impl<T> Deref for RangedValue<T> {
     type Target = T;
 

crates/red_knot_wasm/Cargo.toml

@@ -20,6 +20,7 @@ default = ["console_error_panic_hook"]
 
 [dependencies]
 red_knot_project = { workspace = true, default-features = false, features = ["deflate"] }
+red_knot_ide = { workspace = true }
 red_knot_python_semantic = { workspace = true }
 
 ruff_db = { workspace = true, default-features = false, features = [] }

crates/red_knot_wasm/src/lib.rs

@@ -1,6 +1,7 @@
 use std::any::Any;
 
 use js_sys::{Error, JsString};
+use red_knot_ide::goto_type_definition;
 use red_knot_project::metadata::options::Options;
 use red_knot_project::metadata::value::ValueSource;
 use red_knot_project::watch::{ChangeEvent, ChangedKind, CreatedKind, DeletedKind};
@@ -8,7 +9,7 @@ use red_knot_project::ProjectMetadata;
 use red_knot_project::{Db, ProjectDatabase};
 use red_knot_python_semantic::Program;
 use ruff_db::diagnostic::{self, DisplayDiagnosticConfig};
-use ruff_db::files::{system_path_to_file, File};
+use ruff_db::files::{system_path_to_file, File, FileRange};
 use ruff_db::source::{line_index, source_text};
 use ruff_db::system::walk_directory::WalkDirectoryBuilder;
 use ruff_db::system::{
@@ -17,7 +18,8 @@ use ruff_db::system::{
 };
 use ruff_db::Upcast;
 use ruff_notebook::Notebook;
-use ruff_source_file::SourceLocation;
+use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
+use ruff_text_size::Ranged;
 use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen(start)]
@@ -195,6 +197,52 @@ impl Workspace {
 
         Ok(source_text.to_string())
     }
+
+    #[wasm_bindgen(js_name = "gotoTypeDefinition")]
+    pub fn goto_type_definition(
+        &self,
+        file_id: &FileHandle,
+        position: Position,
+    ) -> Result<Vec<LocationLink>, Error> {
+        let source = source_text(&self.db, file_id.file);
+        let index = line_index(&self.db, file_id.file);
+
+        let offset = index.offset(
+            OneIndexed::new(position.line).ok_or_else(|| {
+                Error::new("Invalid value `0` for `position.line`. The line index is 1-indexed.")
+            })?,
+            OneIndexed::new(position.column).ok_or_else(|| {
+                Error::new(
+                    "Invalid value `0` for `position.column`. The column index is 1-indexed.",
+                )
+            })?,
+            &source,
+        );
+
+        let Some(targets) = goto_type_definition(&self.db, file_id.file, offset) else {
+            return Ok(Vec::new());
+        };
+
+        let source_range = Range::from_text_range(targets.file_range().range(), &index, &source);
+
+        let links: Vec<_> = targets
+            .into_iter()
+            .map(|target| LocationLink {
+                path: target.file().path(&self.db).to_string(),
+                full_range: Range::from_file_range(
+                    &self.db,
+                    FileRange::new(target.file(), target.full_range()),
+                ),
+                selection_range: Some(Range::from_file_range(
+                    &self.db,
+                    FileRange::new(target.file(), target.focus_range()),
+                )),
+                origin_selection_range: Some(source_range),
+            })
+            .collect();
+
+        Ok(links)
+    }
 }
 
 pub(crate) fn into_error<E: std::fmt::Display>(err: E) -> Error {
@@ -257,16 +305,10 @@ impl Diagnostic {
     #[wasm_bindgen(js_name = "toRange")]
     pub fn to_range(&self, workspace: &Workspace) -> Option<Range> {
         self.inner.primary_span().and_then(|span| {
-            let line_index = line_index(workspace.db.upcast(), span.file());
-            let source = source_text(workspace.db.upcast(), span.file());
-            let text_range = span.range()?;
-
-            Some(Range {
-                start: line_index
-                    .source_location(text_range.start(), &source)
-                    .into(),
-                end: line_index.source_location(text_range.end(), &source).into(),
-            })
+            Some(Range::from_file_range(
+                &workspace.db,
+                FileRange::new(span.file(), span.range()?),
+            ))
         })
     }
 
@@ -288,6 +330,38 @@ pub struct Range {
     pub end: Position,
 }
 
+impl Range {
+    fn from_file_range(db: &dyn Db, range: FileRange) -> Self {
+        let index = line_index(db.upcast(), range.file());
+        let source = source_text(db.upcast(), range.file());
+
+        let text_range = range.range();
+
+        let start = index.source_location(text_range.start(), &source);
+        let end = index.source_location(text_range.end(), &source);
+        Self::from((start, end))
+    }
+
+    fn from_text_range(
+        text_range: ruff_text_size::TextRange,
+        line_index: &LineIndex,
+        source: &str,
+    ) -> Self {
+        let start = line_index.source_location(text_range.start(), source);
+        let end = line_index.source_location(text_range.end(), source);
+        Self::from((start, end))
+    }
+}
+
+impl From<(SourceLocation, SourceLocation)> for Range {
+    fn from((start, end): (SourceLocation, SourceLocation)) -> Self {
+        Self {
+            start: start.into(),
+            end: end.into(),
+        }
+    }
+}
+
 #[wasm_bindgen]
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 pub struct Position {
@@ -298,6 +372,14 @@ pub struct Position {
     pub column: usize,
 }
 
+#[wasm_bindgen]
+impl Position {
+    #[wasm_bindgen(constructor)]
+    pub fn new(line: usize, column: usize) -> Self {
+        Self { line, column }
+    }
+}
+
 impl From<SourceLocation> for Position {
     fn from(location: SourceLocation) -> Self {
         Self {
@@ -342,6 +424,20 @@ impl From<ruff_text_size::TextRange> for TextRange {
     }
 }
 
+#[wasm_bindgen]
+pub struct LocationLink {
+    /// The target file path
+    #[wasm_bindgen(getter_with_clone)]
+    pub path: String,
+
+    /// The full range of the target
+    pub full_range: Range,
+    /// The target's range that should be selected/highlighted
+    pub selection_range: Option<Range>,
+    /// The range of the origin.
+    pub origin_selection_range: Option<Range>,
+}
+
 #[derive(Debug, Clone)]
 struct WasmSystem {
     fs: MemoryFileSystem,

playground/knot/src/Editor/Chrome.tsx

@@ -15,7 +15,7 @@ import {
 } from "shared";
 import type { Diagnostic, Workspace } from "red_knot_wasm";
 import { Panel, PanelGroup } from "react-resizable-panels";
-import { Files } from "./Files";
+import { Files, isPythonFile } from "./Files";
 import SecondarySideBar from "./SecondarySideBar";
 import SecondaryPanel, {
   SecondaryPanelResult,
@@ -161,12 +161,14 @@ export default function Chrome({
                   <Editor
                     theme={theme}
                     visible={true}
+                    files={files}
+                    selected={files.selected}
                     fileName={selectedFileName}
-                    source={files.contents[files.selected]}
                     diagnostics={checkResult.diagnostics}
                     workspace={workspace}
                     onMount={handleEditorMount}
                     onChange={(content) => onFileChanged(workspace, content)}
+                    onOpenFile={onFileSelected}
                   />
                   {checkResult.error ? (
                     <div
@@ -245,10 +247,7 @@ function useCheckResult(
     }
 
     const currentHandle = files.handles[files.selected];
-
-    const extension =
-      currentHandle?.path()?.toLowerCase().split(".").pop() ?? "";
-    if (currentHandle == null || !["py", "pyi", "pyw"].includes(extension)) {
+    if (currentHandle == null || !isPythonFile(currentHandle)) {
       return {
         diagnostics: [],
         error: null,

playground/knot/src/Editor/Editor.tsx

@@ -3,49 +3,83 @@
  */
 
 import Moncao, { Monaco, OnMount } from "@monaco-editor/react";
-import { editor, MarkerSeverity } from "monaco-editor";
-import { useCallback, useEffect, useRef } from "react";
+import {
+  CancellationToken,
+  editor,
+  IDisposable,
+  IPosition,
+  IRange,
+  languages,
+  MarkerSeverity,
+  Position,
+  Uri,
+} from "monaco-editor";
+import { RefObject, useCallback, useEffect, useRef } from "react";
 import { Theme } from "shared";
-import { Diagnostic, Severity, Workspace } from "red_knot_wasm";
+import {
+  Diagnostic,
+  Severity,
+  Workspace,
+  Position as KnotPosition,
+  type Range as KnotRange,
+} from "red_knot_wasm";
 
 import IStandaloneCodeEditor = editor.IStandaloneCodeEditor;
+import { FileId, ReadonlyFiles } from "../Playground";
+import { isPythonFile } from "./Files";
 
 type Props = {
   visible: boolean;
   fileName: string;
-  source: string;
+  selected: FileId;
+  files: ReadonlyFiles;
   diagnostics: Diagnostic[];
   theme: Theme;
   workspace: Workspace;
   onChange(content: string): void;
   onMount(editor: IStandaloneCodeEditor, monaco: Monaco): void;
-};
-
-type MonacoEditorState = {
-  monaco: Monaco;
+  onOpenFile(file: FileId): void;
 };
 
 export default function Editor({
   visible,
-  source,
   fileName,
+  selected,
+  files,
   theme,
   diagnostics,
   workspace,
   onChange,
   onMount,
+  onOpenFile,
 }: Props) {
-  const monacoRef = useRef<MonacoEditorState | null>(null);
+  const disposable = useRef<{
+    typeDefinition: IDisposable;
+    editorOpener: IDisposable;
+  } | null>(null);
+  const playgroundState = useRef<PlaygroundServerProps>({
+    monaco: null,
+    files,
+    workspace,
+    onOpenFile,
+  });
+
+  playgroundState.current = {
+    monaco: playgroundState.current.monaco,
+    files,
+    workspace,
+    onOpenFile,
+  };
 
   // Update the diagnostics in the editor.
   useEffect(() => {
-    const editorState = monacoRef.current;
+    const monaco = playgroundState.current.monaco;
 
-    if (editorState == null) {
+    if (monaco == null) {
       return;
     }
 
-    updateMarkers(editorState.monaco, workspace, diagnostics);
+    updateMarkers(monaco, workspace, diagnostics);
   }, [workspace, diagnostics]);
 
   const handleChange = useCallback(
@@ -55,14 +89,30 @@ export default function Editor({
     [onChange],
   );
 
+  useEffect(() => {
+    return () => {
+      disposable.current?.typeDefinition.dispose();
+      disposable.current?.editorOpener.dispose();
+    };
+  }, []);
+
   const handleMount: OnMount = useCallback(
     (editor, instance) => {
       updateMarkers(instance, workspace, diagnostics);
 
-      monacoRef.current = {
-        monaco: instance,
+      const server = new PlaygroundServer(playgroundState);
+      const typeDefinitionDisposable =
+        instance.languages.registerTypeDefinitionProvider("python", server);
+      const editorOpenerDisposable =
+        instance.editor.registerEditorOpener(server);
+
+      disposable.current = {
+        typeDefinition: typeDefinitionDisposable,
+        editorOpener: editorOpenerDisposable,
       };
 
+      playgroundState.current.monaco = instance;
+
       onMount(editor, instance);
     },
 
@@ -79,13 +129,13 @@ export default function Editor({
         fontSize: 14,
         roundedSelection: false,
         scrollBeyondLastLine: false,
-        contextmenu: false,
+        contextmenu: true,
       }}
       language={fileName.endsWith(".pyi") ? "python" : undefined}
       path={fileName}
       wrapperProps={visible ? {} : { style: { display: "none" } }}
       theme={theme === "light" ? "Ayu-Light" : "Ayu-Dark"}
-      value={source}
+      value={files.contents[selected]}
       onChange={handleChange}
     />
   );
@@ -135,3 +185,130 @@ function updateMarkers(
     }),
   );
 }
+
+interface PlaygroundServerProps {
+  monaco: Monaco | null;
+  workspace: Workspace;
+  files: ReadonlyFiles;
+
+  onOpenFile: (file: FileId) => void;
+}
+
+class PlaygroundServer
+  implements languages.TypeDefinitionProvider, editor.ICodeEditorOpener
+{
+  constructor(private props: RefObject<PlaygroundServerProps>) {}
+
+  provideTypeDefinition(
+    model: editor.ITextModel,
+    position: Position,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _: CancellationToken,
+  ): languages.ProviderResult<languages.Definition | languages.LocationLink[]> {
+    const workspace = this.props.current.workspace;
+
+    const selectedFile = this.props.current.files.selected;
+    if (selectedFile == null) {
+      return;
+    }
+
+    const selectedHandle = this.props.current.files.handles[selectedFile];
+
+    if (selectedHandle == null) {
+      return;
+    }
+
+    const links = workspace.gotoTypeDefinition(
+      selectedHandle,
+      new KnotPosition(position.lineNumber, position.column),
+    );
+
+    const locations = links.map((link) => {
+      const targetSelection =
+        link.selection_range == null
+          ? undefined
+          : knotRangeToIRange(link.selection_range);
+
+      const originSelection =
+        link.origin_selection_range == null
+          ? undefined
+          : knotRangeToIRange(link.origin_selection_range);
+
+      return {
+        uri: Uri.parse(link.path),
+        range: knotRangeToIRange(link.full_range),
+        targetSelectionRange: targetSelection,
+        originSelectionRange: originSelection,
+      } as languages.LocationLink;
+    });
+
+    return locations;
+  }
+
+  openCodeEditor(
+    source: editor.ICodeEditor,
+    resource: Uri,
+    selectionOrPosition?: IRange | IPosition,
+  ): boolean {
+    const files = this.props.current.files;
+    const monaco = this.props.current.monaco;
+
+    if (monaco == null) {
+      return false;
+    }
+
+    const fileId = files.index.find((file) => {
+      return Uri.file(file.name).toString() === resource.toString();
+    })?.id;
+
+    if (fileId == null) {
+      return false;
+    }
+
+    const handle = files.handles[fileId];
+
+    let model = monaco.editor.getModel(resource);
+    if (model == null) {
+      const language =
+        handle != null && isPythonFile(handle) ? "python" : undefined;
+      model = monaco.editor.createModel(
+        files.contents[fileId],
+        language,
+        resource,
+      );
+    }
+
+    // it's a bit hacky to create the model manually
+    // but only using `onOpenFile` isn't enough
+    // because the model doesn't get updated until the next render.
+    if (files.selected !== fileId) {
+      source.setModel(model);
+
+      this.props.current.onOpenFile(fileId);
+    }
+
+    if (selectionOrPosition != null) {
+      if (Position.isIPosition(selectionOrPosition)) {
+        source.setPosition(selectionOrPosition);
+        source.revealPosition(selectionOrPosition);
+      } else {
+        source.setSelection(selectionOrPosition);
+        source.revealPosition({
+          lineNumber: selectionOrPosition.startLineNumber,
+          column: selectionOrPosition.startColumn,
+        });
+      }
+    }
+
+    return true;
+  }
+}
+
+function knotRangeToIRange(range: KnotRange): IRange {
+  return {
+    startLineNumber: range.start.line,
+    startColumn: range.start.column,
+    endLineNumber: range.end.line,
+    endColumn: range.end.column,
+  };
+}

playground/knot/src/Editor/Files.tsx

@@ -2,6 +2,7 @@ import { Icons, Theme } from "shared";
 import classNames from "classnames";
 import { useState } from "react";
 import { FileId } from "../Playground";
+import { type FileHandle } from "red_knot_wasm";
 
 export interface Props {
   // The file names
@@ -196,3 +197,8 @@ function FileEntry({ name, onClicked, onRenamed, selected }: FileEntryProps) {
     </button>
   );
 }
+
+export function isPythonFile(handle: FileHandle): boolean {
+  const extension = handle?.path().toLowerCase().split(".").pop() ?? "";
+  return ["py", "pyi", "pyw"].includes(extension);
+}