[ty] Implemented support for "rename" language server feature (#19551)

This PR adds support for the "rename" language server feature. It builds
upon existing functionality used for "go to references".

The "rename" feature involves two language server requests. The first is
a "prepare rename" request that determines whether renaming should be
possible for the identifier at the current offset. The second is a
"rename" request that returns a list of file ranges where the rename
should be applied.

Care must be taken when attempting to rename symbols that span files,
especially if the symbols are defined in files that are not part of the
project. We don't want to modify code in the user's Python environment
or in the vendored stub files.

I found a few bugs in the "go to references" feature when implementing
"rename", and those bug fixes are included in this PR.

---------

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

crates/ty_server/src/capabilities.rs

@@ -1,11 +1,11 @@
 use lsp_types::{
     ClientCapabilities, CompletionOptions, DeclarationCapability, DiagnosticOptions,
     DiagnosticServerCapabilities, HoverProviderCapability, InlayHintOptions,
-    InlayHintServerCapabilities, MarkupKind, OneOf, SelectionRangeProviderCapability,
-    SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions,
-    SemanticTokensServerCapabilities, ServerCapabilities, SignatureHelpOptions,
-    TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
-    TypeDefinitionProviderCapability, WorkDoneProgressOptions,
+    InlayHintServerCapabilities, MarkupKind, OneOf, RenameOptions,
+    SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
+    SemanticTokensOptions, SemanticTokensServerCapabilities, ServerCapabilities,
+    SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
+    TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
 };
 
 use crate::PositionEncoding;
@@ -289,6 +289,10 @@ pub(crate) fn server_capabilities(
         definition_provider: Some(OneOf::Left(true)),
         declaration_provider: Some(DeclarationCapability::Simple(true)),
         references_provider: Some(OneOf::Left(true)),
+        rename_provider: Some(OneOf::Right(RenameOptions {
+            prepare_provider: Some(true),
+            work_done_progress_options: WorkDoneProgressOptions::default(),
+        })),
         document_highlight_provider: Some(OneOf::Left(true)),
         hover_provider: Some(HoverProviderCapability::Simple(true)),
         signature_help_provider: Some(SignatureHelpOptions {

crates/ty_server/src/server/api.rs

@@ -80,6 +80,12 @@ pub(super) fn request(req: server::Request) -> Task {
         requests::SignatureHelpRequestHandler::METHOD => background_document_request_task::<
             requests::SignatureHelpRequestHandler,
         >(req, BackgroundSchedule::Worker),
+        requests::PrepareRenameRequestHandler::METHOD => background_document_request_task::<
+            requests::PrepareRenameRequestHandler,
+        >(req, BackgroundSchedule::Worker),
+        requests::RenameRequestHandler::METHOD => background_document_request_task::<
+            requests::RenameRequestHandler,
+        >(req, BackgroundSchedule::Worker),
         requests::CompletionRequestHandler::METHOD => background_document_request_task::<
             requests::CompletionRequestHandler,
         >(

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

@@ -8,6 +8,8 @@ mod goto_references;
 mod goto_type_definition;
 mod hover;
 mod inlay_hints;
+mod prepare_rename;
+mod rename;
 mod selection_range;
 mod semantic_tokens;
 mod semantic_tokens_range;
@@ -26,6 +28,8 @@ pub(super) use goto_references::ReferencesRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
 pub(super) use inlay_hints::InlayHintRequestHandler;
+pub(super) use prepare_rename::PrepareRenameRequestHandler;
+pub(super) use rename::RenameRequestHandler;
 pub(super) use selection_range::SelectionRangeRequestHandler;
 pub(super) use semantic_tokens::SemanticTokensRequestHandler;
 pub(super) use semantic_tokens_range::SemanticTokensRangeRequestHandler;

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

@@ -0,0 +1,60 @@
+use std::borrow::Cow;
+
+use lsp_types::request::PrepareRenameRequest;
+use lsp_types::{PrepareRenameResponse, TextDocumentPositionParams, Url};
+use ruff_db::source::{line_index, source_text};
+use ty_ide::can_rename;
+use ty_project::ProjectDatabase;
+
+use crate::document::{PositionExt, ToRangeExt};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
+use crate::session::DocumentSnapshot;
+use crate::session::client::Client;
+
+pub(crate) struct PrepareRenameRequestHandler;
+
+impl RequestHandler for PrepareRenameRequestHandler {
+    type RequestType = PrepareRenameRequest;
+}
+
+impl BackgroundDocumentRequestHandler for PrepareRenameRequestHandler {
+    fn document_url(params: &TextDocumentPositionParams) -> Cow<Url> {
+        Cow::Borrowed(&params.text_document.uri)
+    }
+
+    fn run_with_snapshot(
+        db: &ProjectDatabase,
+        snapshot: &DocumentSnapshot,
+        _client: &Client,
+        params: TextDocumentPositionParams,
+    ) -> crate::server::Result<Option<PrepareRenameResponse>> {
+        if snapshot
+            .workspace_settings()
+            .is_language_services_disabled()
+        {
+            return Ok(None);
+        }
+
+        let Some(file) = snapshot.file(db) else {
+            return Ok(None);
+        };
+
+        let source = source_text(db, file);
+        let line_index = line_index(db, file);
+        let offset = params
+            .position
+            .to_text_size(&source, &line_index, snapshot.encoding());
+
+        let Some(range) = can_rename(db, file, offset) else {
+            return Ok(None);
+        };
+
+        let lsp_range = range.to_lsp_range(&source, &line_index, snapshot.encoding());
+
+        Ok(Some(PrepareRenameResponse::Range(lsp_range)))
+    }
+}
+
+impl RetriableRequestHandler for PrepareRenameRequestHandler {}

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

@@ -0,0 +1,83 @@
+use std::borrow::Cow;
+use std::collections::HashMap;
+
+use lsp_types::request::Rename;
+use lsp_types::{RenameParams, TextEdit, Url, WorkspaceEdit};
+use ruff_db::source::{line_index, source_text};
+use ty_ide::rename;
+use ty_project::ProjectDatabase;
+
+use crate::document::{PositionExt, ToLink};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
+use crate::session::DocumentSnapshot;
+use crate::session::client::Client;
+
+pub(crate) struct RenameRequestHandler;
+
+impl RequestHandler for RenameRequestHandler {
+    type RequestType = Rename;
+}
+
+impl BackgroundDocumentRequestHandler for RenameRequestHandler {
+    fn document_url(params: &RenameParams) -> Cow<Url> {
+        Cow::Borrowed(&params.text_document_position.text_document.uri)
+    }
+
+    fn run_with_snapshot(
+        db: &ProjectDatabase,
+        snapshot: &DocumentSnapshot,
+        _client: &Client,
+        params: RenameParams,
+    ) -> crate::server::Result<Option<WorkspaceEdit>> {
+        if snapshot
+            .workspace_settings()
+            .is_language_services_disabled()
+        {
+            return Ok(None);
+        }
+
+        let Some(file) = snapshot.file(db) else {
+            return Ok(None);
+        };
+
+        let source = source_text(db, file);
+        let line_index = line_index(db, file);
+        let offset = params.text_document_position.position.to_text_size(
+            &source,
+            &line_index,
+            snapshot.encoding(),
+        );
+
+        let Some(rename_results) = rename(db, file, offset, &params.new_name) else {
+            return Ok(None);
+        };
+
+        // Group text edits by file
+        let mut changes: HashMap<Url, Vec<TextEdit>> = HashMap::new();
+
+        for reference in rename_results {
+            if let Some(location) = reference.to_location(db, snapshot.encoding()) {
+                let edit = TextEdit {
+                    range: location.range,
+                    new_text: params.new_name.clone(),
+                };
+
+                changes.entry(location.uri).or_default().push(edit);
+            }
+        }
+
+        if changes.is_empty() {
+            return Ok(None);
+        }
+
+        Ok(Some(WorkspaceEdit {
+            changes: Some(changes),
+            document_changes: None,
+            change_annotations: None,
+        }))
+    }
+}
+
+impl RetriableRequestHandler for RenameRequestHandler {}