Add support for local variables in find-references

crates/ark/src/lsp/find_references.rs

@@ -1,16 +1,11 @@
-//
-// references.rs
-//
-// Copyright (C) 2022-2026 Posit Software, PBC. All rights reserved.
-//
-//
-
 use std::path::Path;
 
+use aether_lsp_utils::proto::from_proto;
+use aether_lsp_utils::proto::to_proto;
 use anyhow::anyhow;
 use stdext::result::ResultExt;
+use stdext::unwrap;
 use stdext::unwrap::IntoResult;
-use stdext::*;
 use tower_lsp::lsp_types::Location;
 use tower_lsp::lsp_types::Position;
 use tower_lsp::lsp_types::Range;
@@ -32,6 +27,68 @@ use crate::treesitter::ExtractOperatorType;
 use crate::treesitter::NodeType;
 use crate::treesitter::NodeTypeExt;
 
+pub(crate) fn find_references(
+    params: ReferenceParams,
+    state: &WorldState,
+) -> anyhow::Result<Vec<Location>> {
+    let uri = params.text_document_position.text_document.uri;
+    let position = params.text_document_position.position;
+    let include_declaration = params.context.include_declaration;
+
+    let document = state.get_document(&uri)?;
+
+    let mut locations: Vec<Location> = Vec::new();
+
+    let index = document.semantic_index();
+    let root = document.syntax()?;
+
+    // Intra-file resolution is precise via the semantic index
+    let offset = from_proto::offset_from_position(
+        position,
+        &document.line_index,
+        document.position_encoding,
+    )?;
+    let pos = oak_ide::FileOffset {
+        file: uri.clone(),
+        offset,
+    };
+    let intra = oak_ide::find_references(&index, &root, &pos, include_declaration);
+
+    for file_range in intra {
+        let Some(range) = to_proto::range(
+            file_range.range,
+            &document.line_index,
+            document.position_encoding,
+        )
+        .log_err() else {
+            continue;
+        };
+        locations.push(Location::new(file_range.file, range));
+    }
+
+    if !locations.is_empty() {
+        // If the intra-file pass resolved cleanly, we have the precise answer.
+        // The textual cross-file walk can only add unrefined external matches
+        // which, for a locally-bound symbol, would just be unrelated noise.
+        return Ok(locations);
+    }
+
+    // Truly free variable: no within-file binding. Fall back to a textual walk
+    // over workspace folders (including the current file: intra-file gave
+    // nothing, so no dedup needed). No semantic refinement yet, so unrelated
+    // same-name symbols in other files leak through.
+    if let Ok(context) = build_context(&uri, position, state) {
+        for folder in state.workspace.folders.iter() {
+            if let Ok(path) = folder.to_file_path() {
+                lsp::log_info!("searching references in folder {}", path.display());
+                find_references_in_folder(&context, &path, &mut locations, state);
+            }
+        }
+    }
+
+    Ok(locations)
+}
+
 #[derive(Debug, PartialEq)]
 enum ReferenceKind {
     Symbol, // a regular R symbol
@@ -94,21 +151,19 @@ fn found_match(node: &Node, contents: &str, context: &Context) -> bool {
     if !node.is_identifier() {
         return false;
     }
-
-    let symbol = node.node_to_string(contents).unwrap();
+    let Ok(symbol) = node.node_to_string(contents) else {
+        return false;
+    };
     if symbol != context.symbol {
         return false;
     }
-
     context.kind == node_reference_kind(node)
 }
 
 fn build_context(uri: &Url, position: Position, state: &WorldState) -> anyhow::Result<Context> {
-    // Unwrap the URL.
     let path = uri.file_path()?;
 
-    // Figure out the identifier we're looking for.
-    let context = with_document(path.as_path(), state, |document| {
+    with_document(path.as_path(), state, |document| {
         let ast = &document.ast;
         let contents = document.contents.as_str();
         let point = document.tree_sitter_point_from_lsp_position(position)?;
@@ -118,38 +173,43 @@ fn build_context(uri: &Url, position: Position, state: &WorldState) -> anyhow::R
             .descendant_for_point_range(point, point)
             .into_result()?;
 
-        // Check and see if we got an identifier. If we didn't, we might need to use
-        // some heuristics to look around. Unfortunately, it seems like if you double-click
-        // to select an identifier, and then use Right Click -> Find All References, the
-        // position received by the LSP maps to the _end_ of the selected range, which
-        // is technically not part of the associated identifier's range. In addition, we
-        // can't just subtract 1 from the position column since that would then fail to
-        // resolve the correct identifier when the cursor is located at the start of the
-        // identifier.
+        // Zero-width range queries at an identifier boundary return the
+        // wrapping node rather than the identifier itself. If the cursor is at
+        // the trailing edge of a selection (column past the last character),
+        // retry one column back. If it's at the leading edge (column on the
+        // first character), retry one column forward.
         if !node.is_identifier() && point.column > 0 {
-            let point = Point::new(point.row, point.column - 1);
-            node = ast
+            let back = Point::new(point.row, point.column - 1);
+            if let Some(retry) = ast
                 .root_node()
-                .descendant_for_point_range(point, point)
-                .into_result()?;
+                .descendant_for_point_range(back, back)
+                .filter(|n| n.is_identifier())
+            {
+                node = retry;
+            }
+        }
+        if !node.is_identifier() {
+            let fwd = Point::new(point.row, point.column + 1);
+            if let Some(retry) = ast
+                .root_node()
+                .descendant_for_point_range(fwd, fwd)
+                .filter(|n| n.is_identifier())
+            {
+                node = retry;
+            }
         }
 
-        // double check that we found an identifier
         if !node.is_identifier() {
             return Err(anyhow!(
                 "couldn't find an identifier associated with point {point:?}",
             ));
         }
 
         let kind = node_reference_kind(&node);
-
-        // return identifier text contents
         let symbol = node.node_to_string(contents)?;
 
         Ok(Context { kind, symbol })
-    });
-
-    context
+    })
 }
 
 fn find_references_in_folder(
@@ -167,7 +227,6 @@ fn find_references_in_folder(
             continue;
         }
 
-        lsp::log_info!("found R file {}", path.display());
         let result = with_document(path, state, |document| {
             find_references_in_document(context, path, document, locations)
         });
@@ -201,30 +260,3 @@ fn find_references_in_document(
     });
     Ok(())
 }
-
-pub(crate) fn find_references(
-    params: ReferenceParams,
-    state: &WorldState,
-) -> anyhow::Result<Vec<Location>> {
-    // Create our locations vector.
-    let mut locations: Vec<Location> = Vec::new();
-
-    // Extract relevant parameters.
-    let uri = params.text_document_position.text_document.uri;
-    let position = params.text_document_position.position;
-
-    // Figure out what we're looking for.
-    let context = unwrap!(build_context(&uri, position, state), Err(err) => {
-        return Err(anyhow!("Failed to find build context at position {position:?}: {err:?}"));
-    });
-
-    // Now, start searching through workspace folders for references to that identifier.
-    for folder in state.workspace.folders.iter() {
-        if let Ok(path) = folder.to_file_path() {
-            lsp::log_info!("searching references in folder {}", path.display());
-            find_references_in_folder(&context, &path, &mut locations, state);
-        }
-    }
-
-    Ok(locations)
-}

crates/ark/src/lsp/tests.rs

@@ -1 +1,2 @@
+mod find_references;
 mod goto_definition;