Implement usages search in textDocument/references

fixes #100

commit-id:9ffd3efa

Cargo.toml

@@ -55,6 +55,7 @@ itertools = "0.14.0"
 jod-thread = "0.1.2"
 lsp-server = "0.7.8"
 lsp-types = "=0.95.0"
+memchr = "2.7.4"
 salsa = { package = "rust-analyzer-salsa", version = "0.17.0-pre.6" }
 scarb-metadata = "1.13"
 scarb-proc-macro-server-types = "0.1"

src/ide/navigation/references.rs

@@ -22,11 +22,13 @@ pub fn references(params: ReferenceParams, db: &AnalysisDatabase) -> Option<Vec<
     // For all cases we cover here, definition == declaration.
     let declaration = symbol.definition_location(db);
 
+    // Locations where the searched symbol is used.
+    let usages = symbol.usages(db).collect();
+
     let locations = {
         let declaration = declaration.filter(|_| include_declaration);
 
-        // TODO(mkaput): Implement this.
-        let references = vec![];
+        let references = usages.into_iter().map(|usage| (usage.file, usage.span));
 
         declaration.into_iter().chain(references)
     }

src/lang/db/syntax.rs

@@ -1,16 +1,22 @@
 use cairo_lang_diagnostics::ToOption;
 use cairo_lang_filesystem::ids::FileId;
-use cairo_lang_filesystem::span::TextPosition;
+use cairo_lang_filesystem::span::{TextOffset, TextPosition};
 use cairo_lang_parser::db::ParserGroup;
 use cairo_lang_syntax::node::ast::TerminalIdentifier;
 use cairo_lang_syntax::node::kind::SyntaxKind;
 use cairo_lang_syntax::node::{SyntaxNode, Terminal};
 use cairo_lang_utils::Upcast;
 
 // TODO(mkaput): Make this a real Salsa query group with sensible LRU.
-/// Language server-specific extensions to the syntax group of the Cairo compiler.
+/// LS-specific extensions to the syntax group of the Cairo compiler.
 pub trait LsSyntaxGroup: Upcast<dyn ParserGroup> {
-    /// Finds the most specific [`SyntaxNode`] at the given [`TextPosition`] in file.
+    /// Finds the most specific [`SyntaxNode`] at the given [`TextOffset`] in the file.
+    fn find_syntax_node_at_offset(&self, file: FileId, offset: TextOffset) -> Option<SyntaxNode> {
+        let db = self.upcast();
+        Some(db.file_syntax(file).to_option()?.lookup_offset(db.upcast(), offset))
+    }
+
+    /// Finds the most specific [`SyntaxNode`] at the given [`TextPosition`] in the file.
     fn find_syntax_node_at_position(
         &self,
         file: FileId,
@@ -20,11 +26,11 @@ pub trait LsSyntaxGroup: Upcast<dyn ParserGroup> {
         Some(db.file_syntax(file).to_option()?.lookup_position(db.upcast(), position))
     }
 
-    /// Finds a [`TerminalIdentifier`] at the given [`TextPosition`] in file.
+    /// Finds a [`TerminalIdentifier`] at the given [`TextPosition`] in the file.
     ///
     /// The lookup for identifiers is slightly more sophisticated than just looking for an arbitrary
-    /// syntax node, because identifiers are usually what the user is interested in.
-    /// In case when user position is `ident<caret>()`, while regular syntax node lookup would
+    /// syntax node because identifiers usually are what the user is interested in.
+    /// In case when the user position is `ident<caret>()`, while regular syntax node lookup would
     /// return the left paren, a much better UX would be to correct the lookup to the identifier.
     /// Such corrections are always valid and deterministic, because grammar-wise it is not possible
     /// to have two identifiers/keywords being glued to each other.
@@ -48,10 +54,10 @@ pub trait LsSyntaxGroup: Upcast<dyn ParserGroup> {
         })
     }
 
-    /// Traverse tree in root direction.
+    /// Traverse a tree in the root direction.
     ///
-    /// Finds first node with specified kind.
-    /// Returns it's respective child that is the ancestor of `node`.
+    /// Finds the first node with a specified kind.
+    /// Returns its respective child that is the ancestor of `node`.
     fn first_ancestor_of_kind_respective_child(
         &self,
         mut node: SyntaxNode,

src/lang/inspect/defs.rs

@@ -31,11 +31,13 @@ use smol_str::SmolStr;
 use tracing::error;
 
 use crate::lang::db::{AnalysisDatabase, LsSemanticGroup, LsSyntaxGroup};
+use crate::lang::inspect::usages::FindUsages;
 
 /// Keeps information about the symbol that is being searched for/inspected.
 ///
 /// This is an ephemeral data structure.
 /// Do not store it in any kind of state.
+#[derive(Eq, PartialEq)]
 pub enum SymbolDef {
     Item(ItemDef),
     Variable(VariableDef),
@@ -132,7 +134,6 @@ impl SymbolDef {
     }
 
     /// Gets the name of the symbol.
-    #[expect(unused)]
     pub fn name(&self, db: &AnalysisDatabase) -> SmolStr {
         match self {
             SymbolDef::Item(it) => it.name(db),
@@ -142,9 +143,15 @@ impl SymbolDef {
             SymbolDef::Module(it) => it.name(db),
         }
     }
+
+    /// Starts a find-usages search for this symbol.
+    pub fn usages<'a>(&'a self, db: &'a AnalysisDatabase) -> FindUsages<'a> {
+        FindUsages::new(self, db)
+    }
 }
 
 /// Information about the definition of an item (function, trait, impl, module, etc.).
+#[derive(Eq, PartialEq)]
 pub struct ItemDef {
     /// The [`LookupItemId`] associated with the item.
     lookup_item_id: LookupItemId,
@@ -225,6 +232,7 @@ impl ItemDef {
 }
 
 /// Information about the definition of a variable (local, function parameter).
+#[derive(Eq, PartialEq)]
 pub struct VariableDef {
     name: SmolStr,
     var: Binding,
@@ -355,6 +363,7 @@ impl VariableDef {
 }
 
 /// Information about a struct member.
+#[derive(Eq, PartialEq)]
 pub struct MemberDef {
     member_id: MemberId,
     structure: ItemDef,
@@ -389,6 +398,7 @@ impl MemberDef {
 }
 
 /// Information about the definition of a module.
+#[derive(Eq, PartialEq)]
 pub struct ModuleDef {
     id: ModuleId,
     /// A full path to the parent module if [`ModuleId`] points to a submodule,

src/lang/inspect/mod.rs

@@ -3,3 +3,4 @@
 pub mod crates;
 pub mod defs;
 pub mod methods;
+pub mod usages;

src/lang/inspect/usages.rs

@@ -0,0 +1,125 @@
+use std::collections::HashMap;
+use std::ops::ControlFlow;
+use std::sync::Arc;
+
+use cairo_lang_defs::db::DefsGroup;
+use cairo_lang_filesystem::db::FilesGroup;
+use cairo_lang_filesystem::ids::FileId;
+use cairo_lang_filesystem::span::{TextSpan, TextWidth};
+use cairo_lang_syntax::node::ast::TerminalIdentifier;
+use cairo_lang_syntax::node::{Terminal, TypedStablePtr, TypedSyntaxNode};
+use cairo_lang_utils::Upcast;
+use memchr::memmem::Finder;
+use smol_str::format_smolstr;
+
+use crate::lang::db::{AnalysisDatabase, LsSyntaxGroup};
+use crate::lang::inspect::defs::SymbolDef;
+
+macro_rules! flow {
+    ($expr:expr) => {
+        let ControlFlow::Continue(()) = $expr else {
+            return;
+        };
+    };
+}
+
+// TODO(mkaput): Reject text matches that are not word boundaries.
+// TODO(mkaput): Implement search scopes: for example, variables will never be used in other files.
+// TODO(mkaput): Deal with `crate` keyword.
+/// An implementation of the find-usages functionality.
+///
+/// This algorithm is based on the standard IDE trick:
+/// first, a fast text search to get a superset of matches is performed,
+/// and then each match is checked using a precise goto-definition algorithm.
+pub struct FindUsages<'a> {
+    symbol: &'a SymbolDef,
+    db: &'a AnalysisDatabase,
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct FoundUsage {
+    pub file: FileId,
+    pub span: TextSpan,
+}
+
+impl<'a> FindUsages<'a> {
+    pub(super) fn new(symbol: &'a SymbolDef, db: &'a AnalysisDatabase) -> Self {
+        Self { symbol, db }
+    }
+
+    /// Collects all found usages.
+    pub fn collect(self) -> Vec<FoundUsage> {
+        let mut result = vec![];
+        self.search(&mut |usage| {
+            result.push(usage);
+            ControlFlow::Continue(())
+        });
+        result
+    }
+
+    /// Executes this search and calls the given sink for each found usage.
+    #[tracing::instrument(skip_all)]
+    pub fn search(self, sink: &mut dyn FnMut(FoundUsage) -> ControlFlow<(), ()>) {
+        let db = self.db;
+
+        let needle = match self.symbol {
+            // Small optimisation for inline macros: we can be sure that any usages will have a `!`
+            // at the end, so we do not need to search for occurrences without it.
+            SymbolDef::ExprInlineMacro(macro_name) => format_smolstr!("{macro_name}!"),
+            symbol => symbol.name(db),
+        };
+
+        let finder = Finder::new(needle.as_bytes());
+
+        for (file, text) in Self::scope_files(db) {
+            // Search occurrences of the symbol's name.
+            for offset in finder.find_iter(text.as_bytes()) {
+                let offset = TextWidth::at(&text, offset).as_offset();
+                if let Some(node) = db.find_syntax_node_at_offset(file, offset) {
+                    if let Some(identifier) = TerminalIdentifier::cast_token(db.upcast(), node) {
+                        flow!(self.found_identifier(identifier, sink));
+                    }
+                }
+            }
+        }
+    }
+
+    fn scope_files(db: &AnalysisDatabase) -> impl Iterator<Item = (FileId, Arc<str>)> + '_ {
+        let mut files = HashMap::new();
+        for crate_id in db.crates() {
+            for &module_id in db.crate_modules(crate_id).iter() {
+                if let Ok(file_id) = db.module_main_file(module_id) {
+                    if let Some(text) = db.file_content(file_id) {
+                        files.insert(file_id, text);
+                    }
+                }
+            }
+        }
+        files.into_iter()
+    }
+
+    fn found_identifier(
+        &self,
+        identifier: TerminalIdentifier,
+        sink: &mut dyn FnMut(FoundUsage) -> ControlFlow<(), ()>,
+    ) -> ControlFlow<(), ()> {
+        // Declaration is not a usage, so filter it out.
+        if Some(identifier.stable_ptr().untyped()) == self.symbol.definition_stable_ptr() {
+            return ControlFlow::Continue(());
+        }
+        let Some(found_symbol) = SymbolDef::find(self.db, &identifier) else {
+            return ControlFlow::Continue(());
+        };
+        if found_symbol == *self.symbol {
+            let syntax_db = self.db.upcast();
+            let syntax_node = identifier.as_syntax_node();
+            let usage = FoundUsage {
+                file: syntax_node.stable_ptr().file_id(syntax_db),
+                span: syntax_node.span_without_trivia(syntax_db),
+            };
+            sink(usage)
+        } else {
+            ControlFlow::Continue(())
+        }
+    }
+}

tests/test_data/references/fns.txt

@@ -11,6 +11,7 @@ hello = "src"
 edition = "2024_07"
 
 //! > cairo_code
+// FIXME(#129): definition stable ptr for functions is wrong and causes declaration slipping here.
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
 
 fn main() {
@@ -20,10 +21,20 @@ fn main() {
 //! > References #0
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
 ---
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let x = <sel>pow2</sel>(2) + pow2(3);
+
+    let x = pow2(2) + <sel>pow2</sel>(3);
 
 //! > References #1
     let x = po<caret>w2(2) + pow2(3);
 ---
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let x = <sel>pow2</sel>(2) + pow2(3);
+
+    let x = pow2(2) + <sel>pow2</sel>(3);
 
 //! > ==========================================================================
 
@@ -51,11 +62,23 @@ fn pow<caret>2(x: felt252) -> felt252 { x * x }
 ---
 <sel>fn pow2(x: felt252) -> felt252 { x * x }</sel>
 
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let x = <sel>pow2</sel>(2) + pow2(3);
+
+    let x = pow2(2) + <sel>pow2</sel>(3);
+
 //! > References #1
     let x = po<caret>w2(2) + pow2(3);
 ---
 <sel>fn pow2(x: felt252) -> felt252 { x * x }</sel>
 
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let x = <sel>pow2</sel>(2) + pow2(3);
+
+    let x = pow2(2) + <sel>pow2</sel>(3);
+
 //! > ==========================================================================
 
 //! > Test unused function.
@@ -74,13 +97,16 @@ edition = "2024_07"
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
 
 fn main() {
-    let pow2 = 2;  // bad
+    let pow2 = 2;  // bad  // FIXME(mkaput): Why this?
     let x = pow2 + pow2; // bad
 }
 
 //! > References #0
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
 ---
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let <sel>pow2</sel> = 2;  // bad  // FIXME(mkaput): Why this?
 
 //! > ==========================================================================
 
@@ -106,3 +132,8 @@ fn main() {
 //! > References #0
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
 ---
+fn <sel>pow2</sel>(x: felt252) -> felt252 { x * x }
+
+    let x = <sel>pow2</sel>(2) + pow2(3);
+
+    let x = pow2(2) + <sel>pow2</sel>(3);