find-references on reassignments

pyrefly/lib/state/lsp.rs

@@ -1710,11 +1710,51 @@ impl<'a> Transaction<'a> {
                 && definition_range == definition_export.location
             {
                 references.push(key.range());
+            } else if let Key::Definition(id) = key {
+                // Check if this is a different definition of the same variable name
+                if module.code_at(id.range()).trim() == module.code_at(definition_range).trim()
+                    && id.range() != definition_range
+                {
+                    // Only consider it a reference if it's in the same scope level as the original definition
+                    if self.are_in_same_scope(&handle, definition_range, id.range()) {
+                        references.push(id.range());
+                    }
+                }
             }
         }
+
         Some(references)
     }
 
+    fn are_in_same_scope(&self, handle: &Handle, range1: TextRange, range2: TextRange) -> bool {
+        let Some(ast) = self.get_ast(handle) else {
+            return false;
+        };
+
+        let scope1 = Self::find_scope_node(&ast, range1.start());
+        let scope2 = Self::find_scope_node(&ast, range2.start());
+
+        scope1 == scope2
+    }
+
+    fn find_scope_node(ast: &ModModule, position: TextSize) -> Option<TextRange> {
+        let covering_nodes = Ast::locate_node(ast, position);
+
+        for node in covering_nodes.iter() {
+            match node {
+                AnyNodeRef::StmtFunctionDef(func_def) => {
+                    return Some(func_def.range());
+                }
+                AnyNodeRef::StmtClassDef(class_def) => {
+                    return Some(class_def.range());
+                }
+                _ => continue,
+            }
+        }
+
+        None
+    }
+
     /// Bindings can contain synthetic bindings, which are not meaningful to end users.
     /// This function helps to filter out such bindings and only leave bindings that eventually
     /// jumps to a name in the source.

pyrefly/lib/test/lsp/local_find_refs.rs

@@ -299,7 +299,6 @@ References:
     );
 }
 
-// todo(kylei): references on renames should find lhs of assignment
 #[test]
 fn reassigned_local() {
     let code = r#"
@@ -316,10 +315,65 @@ xy = xy + 1
 References:
 2 | xy = 5
     ^^
+4 | xy = xy + 1
+    ^^
 4 | xy = xy + 1
          ^^
 "#
         .trim(),
         report.trim(),
     );
 }
+
+#[test]
+fn reassigned_local_out_of_scope() {
+    let code = r#"
+xy = 5
+#^
+def increment_xy():
+    xy = xy + 1
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+2 | xy = 5
+     ^
+References:
+2 | xy = 5
+    ^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn reassigned_global() {
+    let code = r#"
+xy = 5
+#^
+def increment_xy():
+    global xy
+    xy = xy + 1
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+2 | xy = 5
+     ^
+References:
+2 | xy = 5
+    ^^
+5 |     global xy
+               ^^
+6 |     xy = xy + 1
+        ^^
+6 |     xy = xy + 1
+             ^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}