[ty] Add initial implementation of goto definition for loads of local names

crates/ty_ide/src/lib.rs

@@ -8,7 +8,7 @@ mod markup;
 
 pub use completion::completion;
 pub use db::Db;
-pub use goto::goto_type_definition;
+pub use goto::{goto_definition, goto_type_definition};
 pub use hover::hover;
 pub use inlay_hints::inlay_hints;
 pub use markup::MarkupKind;

crates/ty_python_semantic/src/lib.rs

@@ -15,7 +15,7 @@ pub use program::{
     PythonVersionWithSource, SearchPathSettings,
 };
 pub use python_platform::PythonPlatform;
-pub use semantic_model::{Completion, HasType, NameKind, SemanticModel};
+pub use semantic_model::{Completion, HasDefinition, HasType, NameKind, SemanticModel};
 pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin};
 pub use util::diagnostics::add_inferred_python_version_hint_to_diagnostic;
 

crates/ty_python_semantic/src/semantic_model.rs

@@ -1,12 +1,14 @@
 use ruff_db::files::{File, FilePath};
 use ruff_db::source::line_index;
-use ruff_python_ast as ast;
+use ruff_python_ast::{self as ast, ExprContext};
 use ruff_python_ast::{Expr, ExprRef, name::Name};
 use ruff_source_file::LineIndex;
 
 use crate::Db;
 use crate::module_name::ModuleName;
 use crate::module_resolver::{KnownModule, Module, resolve_module};
+use crate::semantic_index::ast_ids::HasScopedUseId;
+use crate::semantic_index::definition::Definition;
 use crate::semantic_index::place::FileScopeId;
 use crate::semantic_index::semantic_index;
 use crate::types::ide_support::all_declarations_and_bindings;
@@ -175,6 +177,41 @@ pub struct Completion {
     pub builtin: bool,
 }
 
+pub trait HasDefinition {
+    /// Returns the definitions of `self`.
+    ///
+    /// ## Panics
+    /// May panic if `self` is from another file than `model`.
+    fn definitions<'db>(&self, model: &SemanticModel<'db>) -> Option<Vec<Definition<'db>>>;
+}
+
+impl HasDefinition for ast::ExprRef<'_> {
+    fn definitions<'db>(&self, model: &SemanticModel<'db>) -> Option<Vec<Definition<'db>>> {
+        match self {
+            ExprRef::Name(name) => match name.ctx {
+                ExprContext::Load => {
+                    let index = semantic_index(model.db, model.file);
+                    let file_scope = index.expression_scope_id(*self);
+                    let scope = file_scope.to_scope_id(model.db, model.file);
+                    let use_def = index.use_def_map(file_scope);
+                    let use_id = self.scoped_use_id(model.db, scope);
+
+                    Some(
+                        use_def
+                            .bindings_at_use(use_id)
+                            .filter_map(|binding| binding.binding.definition())
+                            .collect(),
+                    )
+                }
+                ExprContext::Store => None,
+                ExprContext::Del => None,
+                ExprContext::Invalid => None,
+            },
+            _ => None,
+        }
+    }
+}
+
 pub trait HasType {
     /// Returns the inferred type of `self`.
     ///

crates/ty_server/src/server.rs

@@ -6,7 +6,7 @@ use crate::session::{AllOptions, ClientOptions, Session};
 use lsp_server::Connection;
 use lsp_types::{
     ClientCapabilities, DiagnosticOptions, DiagnosticServerCapabilities, HoverProviderCapability,
-    InlayHintOptions, InlayHintServerCapabilities, MessageType, ServerCapabilities,
+    InlayHintOptions, InlayHintServerCapabilities, MessageType, OneOf, ServerCapabilities,
     TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
     TypeDefinitionProviderCapability, Url,
 };
@@ -183,6 +183,7 @@ impl Server {
                     ..Default::default()
                 },
             )),
+            definition_provider: Some(OneOf::Left(true)),
             type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
             hover_provider: Some(HoverProviderCapability::Simple(true)),
             inlay_hint_provider: Some(lsp_types::OneOf::Right(

crates/ty_server/src/server/api.rs

@@ -43,6 +43,9 @@ pub(super) fn request(req: server::Request) -> Task {
         >(
             req, BackgroundSchedule::Worker
         ),
+        requests::GotoDefinitionRequestHandler::METHOD => background_document_request_task::<
+            requests::GotoDefinitionRequestHandler,
+        >(req, BackgroundSchedule::Worker),
         requests::HoverRequestHandler::METHOD => background_document_request_task::<
             requests::HoverRequestHandler,
         >(req, BackgroundSchedule::Worker),

crates/ty_server/src/server/api/requests.rs

@@ -1,5 +1,6 @@
 mod completion;
 mod diagnostic;
+mod goto_definition;
 mod goto_type_definition;
 mod hover;
 mod inlay_hints;
@@ -8,6 +9,7 @@ mod workspace_diagnostic;
 
 pub(super) use completion::CompletionRequestHandler;
 pub(super) use diagnostic::DocumentDiagnosticRequestHandler;
+pub(super) use goto_definition::GotoDefinitionRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
 pub(super) use inlay_hints::InlayHintRequestHandler;

crates/ty_server/src/server/api/requests/goto_definition.rs

@@ -0,0 +1,77 @@
+use std::borrow::Cow;
+
+use lsp_types::GotoDefinitionParams;
+use lsp_types::request::GotoDefinition;
+use lsp_types::{GotoDefinitionResponse, Url};
+use ruff_db::source::{line_index, source_text};
+use ty_ide::goto_definition;
+use ty_project::ProjectDatabase;
+
+use crate::DocumentSnapshot;
+use crate::document::{PositionExt, ToLink};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
+use crate::session::client::Client;
+
+pub(crate) struct GotoDefinitionRequestHandler;
+
+impl RequestHandler for GotoDefinitionRequestHandler {
+    type RequestType = GotoDefinition;
+}
+
+impl BackgroundDocumentRequestHandler for GotoDefinitionRequestHandler {
+    fn document_url(params: &GotoDefinitionParams) -> Cow<Url> {
+        Cow::Borrowed(&params.text_document_position_params.text_document.uri)
+    }
+
+    fn run_with_snapshot(
+        db: &ProjectDatabase,
+        snapshot: DocumentSnapshot,
+        _client: &Client,
+        params: GotoDefinitionParams,
+    ) -> crate::server::Result<Option<GotoDefinitionResponse>> {
+        if snapshot.client_settings().is_language_services_disabled() {
+            return Ok(None);
+        }
+
+        let Some(file) = snapshot.file(db) else {
+            tracing::debug!("Failed to resolve file for {:?}", params);
+            return Ok(None);
+        };
+
+        let source = source_text(db, file);
+        let line_index = line_index(db, file);
+        let offset = params.text_document_position_params.position.to_text_size(
+            &source,
+            &line_index,
+            snapshot.encoding(),
+        );
+
+        let Some(ranged) = goto_definition(db, file, offset) else {
+            return Ok(None);
+        };
+
+        if snapshot
+            .resolved_client_capabilities()
+            .type_definition_link_support
+        {
+            let src = Some(ranged.range);
+            let links: Vec<_> = ranged
+                .into_iter()
+                .filter_map(|target| target.to_link(db, src, snapshot.encoding()))
+                .collect();
+
+            Ok(Some(GotoDefinitionResponse::Link(links)))
+        } else {
+            let locations: Vec<_> = ranged
+                .into_iter()
+                .filter_map(|target| target.to_location(db, snapshot.encoding()))
+                .collect();
+
+            Ok(Some(GotoDefinitionResponse::Array(locations)))
+        }
+    }
+}
+
+impl RetriableRequestHandler for GotoDefinitionRequestHandler {}