feat: implement Go to Assignments for local variables

Author: sgasho

crates/ide/src/goto_assignments.rs

@@ -0,0 +1,244 @@
+use hir::Semantics;
+use ide_db::{
+    RootDatabase,
+    defs::{Definition, IdentClass},
+    helpers::pick_best_token,
+    search::ReferenceCategory,
+};
+use syntax::{AstNode, SyntaxKind::IDENT};
+
+use crate::{FilePosition, NavigationTarget, RangeInfo, TryToNav};
+
+/// Find all assignment sites for the variable at the given position.
+///
+/// Returns all locations where the variable is assigned a value, including:
+/// - Initial definition sites (let bindings, function parameters, etc.)
+/// - Explicit assignment expressions (x = value)
+/// - Compound assignment expressions (x += value, x *= value, etc.)
+pub(crate) fn goto_assignments(
+    db: &RootDatabase,
+    position: FilePosition,
+) -> Option<RangeInfo<Vec<NavigationTarget>>> {
+    let sema = &Semantics::new(db);
+
+    let def = find_definition_at_position(sema, position)?;
+
+    let Definition::Local(_) = def else {
+        return None;
+    };
+
+    find_assignments_for_def(sema, def, position)
+}
+
+fn find_definition_at_position(
+    sema: &Semantics<'_, RootDatabase>,
+    position: FilePosition,
+) -> Option<Definition> {
+    let file = sema.parse_guess_edition(position.file_id);
+    let token =
+        pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
+            IDENT => 1,
+            _ => 0,
+        })?;
+
+    let token = sema.descend_into_macros_no_opaque(token, false).pop()?;
+    let parent = token.value.parent()?;
+
+    IdentClass::classify_node(sema, &parent)?.definitions().pop().map(|(def, _)| def)
+}
+
+fn find_assignments_for_def(
+    sema: &Semantics<'_, RootDatabase>,
+    def: Definition,
+    position: FilePosition,
+) -> Option<RangeInfo<Vec<NavigationTarget>>> {
+    let mut targets = Vec::new();
+
+    if let Some(nav_result) = def.try_to_nav(sema.db) {
+        targets.push(nav_result.call_site);
+    }
+
+    let usages = def.usages(sema).include_self_refs().all();
+
+    targets.extend(usages.iter().flat_map(|(file_id, refs)| {
+        refs.iter().filter_map(move |file_ref| {
+            file_ref.category.contains(ReferenceCategory::WRITE).then(|| {
+                NavigationTarget::from_syntax(
+                    file_id.file_id(sema.db),
+                    "assignment".into(),
+                    Some(file_ref.range),
+                    file_ref.range,
+                    ide_db::SymbolKind::Local,
+                )
+            })
+        })
+    }));
+
+    if targets.is_empty() {
+        return None;
+    }
+
+    let range = sema
+        .parse_guess_edition(position.file_id)
+        .syntax()
+        .token_at_offset(position.offset)
+        .next()
+        .map(|token| token.text_range())?;
+
+    Some(RangeInfo::new(range, targets))
+}
+
+#[cfg(test)]
+mod tests {
+    use ide_db::FileRange;
+    use itertools::Itertools;
+
+    use crate::fixture;
+
+    fn check(#[rust_analyzer::rust_fixture] ra_fixture: &str) {
+        let (analysis, position, expected) = fixture::annotations(ra_fixture);
+        let navs = analysis.goto_assignments(position).unwrap().expect("no assignments found").info;
+        if navs.is_empty() {
+            panic!("unresolved reference")
+        }
+
+        let cmp = |&FileRange { file_id, range }: &_| (file_id, range.start());
+        let navs = navs
+            .into_iter()
+            .map(|nav| FileRange { file_id: nav.file_id, range: nav.focus_or_full_range() })
+            .sorted_by_key(cmp)
+            .collect::<Vec<_>>();
+        let expected = expected
+            .into_iter()
+            .map(|(FileRange { file_id, range }, _)| FileRange { file_id, range })
+            .sorted_by_key(cmp)
+            .collect::<Vec<_>>();
+        assert_eq!(expected, navs);
+    }
+
+    #[test]
+    fn goto_assignments_reassignments() {
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    let mut a = 0;
+    let mut x = 1;
+         // ^
+    x$0 = 2;
+ // ^  
+    println!("{}", x);
+    x = 3;
+ // ^   
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn goto_assignments_compound_operators() {
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    let mut x = 10;
+         // ^
+    x += 5;
+ // ^
+    x$0 *= 2;
+ // ^
+    println!("{}", x);
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn goto_assignments_struct_field_mutation() {
+        check(
+            r#"
+//- /main.rs
+struct Point { x: i32, y: i32 }
+
+fn main() {
+    let mut p = Point { x: 0, y: 0 };
+         // ^
+    p$0 = Point { x: 10, y: 20 };
+ // ^
+    p.x = 5;  // This is not an assignment to `p` itself
+    println!("{:?}", p);
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn goto_assignments_immutable_variable() {
+        // Immutable variables only have the initial definition, no assignments
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    let x$0 = 5;
+     // ^
+    println!("{}", x);
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn goto_assignments_closure_capture() {
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    let mut x = 0;
+         // ^
+    let closure = |mut y: i32| {
+        x$0 = 42;
+     // ^
+        y = 1;  // This is a different variable
+    };
+    closure(0);
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn goto_assignments_loop_variable() {
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    for mut i in 0..3 {
+         // ^
+        if i > 1 {
+            i$0 += 1;
+         // ^
+        }
+    }
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn goto_assignments_shadowing() {
+        // Each `x` is a separate variable, so only assignments to the same binding
+        check(
+            r#"
+//- /main.rs
+fn main() {
+    let mut x = 1;
+    let mut x = 2;  // Different variable (shadowing)
+         // ^
+    x$0 = 3;
+ // ^
+    println!("{}", x);
+}
+"#,
+        );
+    }
+}

crates/ide/src/lib.rs

@@ -27,6 +27,7 @@ mod extend_selection;
 mod fetch_crates;
 mod file_structure;
 mod folding_ranges;
+mod goto_assignments;
 mod goto_declaration;
 mod goto_definition;
 mod goto_implementation;
@@ -518,6 +519,14 @@ impl Analysis {
         self.with_db(|db| goto_type_definition::goto_type_definition(db, position))
     }
 
+    /// Returns the type definitions for the symbol at `position`.
+    pub fn goto_assignments(
+        &self,
+        position: FilePosition,
+    ) -> Cancellable<Option<RangeInfo<Vec<NavigationTarget>>>> {
+        self.with_db(|db| goto_assignments::goto_assignments(db, position))
+    }
+
     pub fn find_all_refs(
         &self,
         position: FilePosition,

crates/rust-analyzer/src/handlers/request.rs

@@ -32,6 +32,7 @@ use syntax::{TextRange, TextSize};
 use triomphe::Arc;
 use vfs::{AbsPath, AbsPathBuf, FileId, VfsPath};
 
+use crate::lsp::ext::GotoAssignmentsResponse;
 use crate::{
     config::{Config, RustfmtConfig, WorkspaceSymbolConfig},
     diagnostics::convert_diagnostic,
@@ -870,6 +871,21 @@ pub(crate) fn handle_goto_type_definition(
     Ok(Some(res))
 }
 
+pub(crate) fn handle_goto_assignments(
+    snap: GlobalStateSnapshot,
+    params: lsp_types::TextDocumentPositionParams,
+) -> anyhow::Result<Option<GotoAssignmentsResponse>> {
+    let _p = tracing::info_span!("handle_goto_assignments").entered();
+    let position = try_default!(from_proto::file_position(&snap, params)?);
+    let nav_info = match snap.analysis.goto_assignments(position)? {
+        None => return Ok(None),
+        Some(it) => it,
+    };
+    let src = FileRange { file_id: position.file_id, range: nav_info.range };
+    let res = to_proto::goto_assignments_response(&snap, Some(src), nav_info.info)?;
+    Ok(Some(res))
+}
+
 pub(crate) fn handle_parent_module(
     snap: GlobalStateSnapshot,
     params: lsp_types::TextDocumentPositionParams,

crates/rust-analyzer/src/lsp/ext.rs

@@ -407,6 +407,40 @@ impl Request for ChildModules {
     const METHOD: &'static str = "experimental/childModules";
 }
 
+pub enum GotoAssignments {}
+
+#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
+#[serde(untagged)]
+pub enum GotoAssignmentsResponse {
+    Scalar(lsp_types::Location),
+    Array(Vec<lsp_types::Location>),
+    Link(Vec<lsp_types::LocationLink>),
+}
+
+impl From<lsp_types::Location> for GotoAssignmentsResponse {
+    fn from(location: lsp_types::Location) -> Self {
+        GotoAssignmentsResponse::Scalar(location)
+    }
+}
+
+impl From<Vec<lsp_types::Location>> for GotoAssignmentsResponse {
+    fn from(locations: Vec<lsp_types::Location>) -> Self {
+        GotoAssignmentsResponse::Array(locations)
+    }
+}
+
+impl From<Vec<lsp_types::LocationLink>> for GotoAssignmentsResponse {
+    fn from(locations: Vec<lsp_types::LocationLink>) -> Self {
+        GotoAssignmentsResponse::Link(locations)
+    }
+}
+
+impl Request for GotoAssignments {
+    type Params = lsp_types::TextDocumentPositionParams;
+    type Result = Option<GotoAssignmentsResponse>;
+    const METHOD: &'static str = "experimental/gotoAssignments";
+}
+
 pub enum JoinLines {}
 
 impl Request for JoinLines {

crates/rust-analyzer/src/lsp/to_proto.rs

@@ -23,6 +23,7 @@ use semver::VersionReq;
 use serde_json::to_value;
 use vfs::AbsPath;
 
+use crate::lsp::ext::GotoAssignmentsResponse;
 use crate::{
     config::{CallInfoConfig, Config},
     global_state::GlobalStateSnapshot,
@@ -1080,6 +1081,29 @@ pub(crate) fn goto_definition_response(
     }
 }
 
+pub(crate) fn goto_assignments_response(
+    snap: &GlobalStateSnapshot,
+    src: Option<FileRange>,
+    targets: Vec<NavigationTarget>,
+) -> Cancellable<GotoAssignmentsResponse> {
+    if snap.config.location_link() {
+        let links = targets
+            .into_iter()
+            .unique_by(|nav| (nav.file_id, nav.full_range, nav.focus_range))
+            .map(|nav| location_link(snap, src, nav))
+            .collect::<Cancellable<Vec<_>>>()?;
+        Ok(links.into())
+    } else {
+        let locations = targets
+            .into_iter()
+            .map(|nav| FileRange { file_id: nav.file_id, range: nav.focus_or_full_range() })
+            .unique()
+            .map(|range| location(snap, range))
+            .collect::<Cancellable<Vec<_>>>()?;
+        Ok(locations.into())
+    }
+}
+
 fn outside_workspace_annotation_id() -> String {
     String::from("OutsideWorkspace")
 }

crates/rust-analyzer/src/main_loop.rs

@@ -1202,6 +1202,7 @@ impl GlobalState {
             .on::<NO_RETRY, lsp_ext::ExpandMacro>(handlers::handle_expand_macro)
             .on::<NO_RETRY, lsp_ext::ParentModule>(handlers::handle_parent_module)
             .on::<NO_RETRY, lsp_ext::ChildModules>(handlers::handle_child_modules)
+            .on::<NO_RETRY, lsp_ext::GotoAssignments>(handlers::handle_goto_assignments)
             .on::<NO_RETRY, lsp_ext::Runnables>(handlers::handle_runnables)
             .on::<NO_RETRY, lsp_ext::RelatedTests>(handlers::handle_related_tests)
             .on::<NO_RETRY, lsp_ext::CodeActionRequest>(handlers::handle_code_action)