From 64fe6e14d30121869493940a1e22f6f8d5a089df Mon Sep 17 00:00:00 2001
From: Marek Kaput <marek.kaput@swmansion.com>
Date: Fri, 20 Dec 2024 13:33:28 +0100
Subject: [PATCH] Collect declarations in `textDocument/references`
 implementation

commit-id:67cec2e0
---
 src/ide/navigation/references.rs   | 39 +++++++++++++++++++++++++++---
 src/lang/inspect/mod.rs            |  1 +
 src/lang/inspect/refs.rs           | 36 +++++++++++++++++++++++++++
 tests/test_data/references/fns.txt | 10 +++++---
 4 files changed, 78 insertions(+), 8 deletions(-)
 create mode 100644 src/lang/inspect/refs.rs

diff --git a/src/ide/navigation/references.rs b/src/ide/navigation/references.rs
index 316b5a0c..d134c7bd 100644
--- a/src/ide/navigation/references.rs
+++ b/src/ide/navigation/references.rs
@@ -1,8 +1,39 @@
+use cairo_lang_utils::Upcast;
+use itertools::Itertools;
 use lsp_types::{Location, ReferenceParams};
 
-use crate::lang::db::AnalysisDatabase;
+use crate::lang::db::{AnalysisDatabase, LsSyntaxGroup};
+use crate::lang::inspect::refs::find_all_references;
+use crate::lang::lsp::{LsProtoGroup, ToCairo, ToLsp};
 
-pub fn references(_params: ReferenceParams, _db: &AnalysisDatabase) -> Option<Vec<Location>> {
-    // TODO(mkaput): Implement this.
-    None
+pub fn references(params: ReferenceParams, db: &AnalysisDatabase) -> Option<Vec<Location>> {
+    let include_declaration = params.context.include_declaration;
+
+    let file = db.file_for_url(&params.text_document_position.text_document.uri)?;
+    let position = params.text_document_position.position.to_cairo();
+
+    let identifier = db.find_identifier_at_position(file, position)?;
+
+    let found_references = find_all_references(db, identifier)?;
+
+    let locations = {
+        let declaration = found_references.declaration.filter(|_| include_declaration);
+
+        let references = found_references
+            .usages
+            .into_iter()
+            .flat_map(|(file, spans)| spans.into_iter().map(move |span| (file, span)));
+
+        declaration.into_iter().chain(references)
+    }
+    .unique()
+    .filter_map(|(file, span)| {
+        let found_uri = db.url_for_file(file)?;
+        let range = span.position_in_file(db.upcast(), file)?.to_lsp();
+        let location = Location { uri: found_uri, range };
+        Some(location)
+    })
+    .collect();
+
+    Some(locations)
 }
diff --git a/src/lang/inspect/mod.rs b/src/lang/inspect/mod.rs
index 308a8b12..8f428eeb 100644
--- a/src/lang/inspect/mod.rs
+++ b/src/lang/inspect/mod.rs
@@ -3,3 +3,4 @@
 pub mod crates;
 pub mod defs;
 pub mod methods;
+pub mod refs;
diff --git a/src/lang/inspect/refs.rs b/src/lang/inspect/refs.rs
new file mode 100644
index 00000000..a0bb9cc0
--- /dev/null
+++ b/src/lang/inspect/refs.rs
@@ -0,0 +1,36 @@
+use std::collections::HashMap;
+
+use cairo_lang_filesystem::ids::FileId;
+use cairo_lang_filesystem::span::TextSpan;
+use cairo_lang_syntax::node::ast::TerminalIdentifier;
+
+use crate::lang::db::AnalysisDatabase;
+use crate::lang::inspect::defs::SymbolDef;
+
+pub struct FoundReferences {
+    /// Location where searched symbol is declared.
+    ///
+    /// This can rarely be `None`, for example, for macros.
+    pub declaration: Option<(FileId, TextSpan)>,
+
+    /// Locations where searched symbol is used.
+    pub usages: HashMap<FileId, Vec<TextSpan>>,
+}
+
+/// Finds all places in the entire analysed codebase for usages of the given identifier.
+pub fn find_all_references(
+    db: &AnalysisDatabase,
+    identifier: TerminalIdentifier,
+) -> Option<FoundReferences> {
+    let symbol = SymbolDef::find(db, &identifier)?;
+
+    // TODO(mkaput): Think about how to deal with `mod foo;` vs `mod foo { ... }`.
+    // For all cases we cover here, definition == declaration.
+    let declaration = symbol.definition_location(db);
+
+    Some(FoundReferences {
+        declaration,
+        // TODO(mkaput): Implement this.
+        usages: Default::default(),
+    })
+}
diff --git a/tests/test_data/references/fns.txt b/tests/test_data/references/fns.txt
index 28a42b61..3e209e96 100644
--- a/tests/test_data/references/fns.txt
+++ b/tests/test_data/references/fns.txt
@@ -19,11 +19,11 @@ fn main() {
 
 //! > References #0
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
-none response
+---
 
 //! > References #1
     let x = po<caret>w2(2) + pow2(3);
-none response
+---
 
 //! > ==========================================================================
 
@@ -48,8 +48,10 @@ fn main() {
 
 //! > References #0
 fn pow<caret>2(x: felt252) -> felt252 { x * x }
-none response
+---
+<sel>fn pow2(x: felt252) -> felt252 { x * x }</sel>
 
 //! > References #1
     let x = po<caret>w2(2) + pow2(3);
-none response
+---
+<sel>fn pow2(x: felt252) -> felt252 { x * x }</sel>