1use language::BufferSnapshot;
  2use std::collections::HashMap;
  3use std::ops::Range;
  4use util::RangeExt;
  5
  6use crate::{
  7    declaration::Identifier,
  8    excerpt::{EditPredictionExcerpt, EditPredictionExcerptText},
  9};
 10
 11#[derive(Debug)]
 12pub struct Reference {
 13    pub identifier: Identifier,
 14    pub range: Range<usize>,
 15    pub region: ReferenceRegion,
 16}
 17
 18#[derive(Copy, Clone, Debug, Eq, PartialEq)]
 19pub enum ReferenceRegion {
 20    Breadcrumb,
 21    Nearby,
 22}
 23
 24pub fn references_in_excerpt(
 25    excerpt: &EditPredictionExcerpt,
 26    excerpt_text: &EditPredictionExcerptText,
 27    snapshot: &BufferSnapshot,
 28) -> HashMap<Identifier, Vec<Reference>> {
 29    let mut references = identifiers_in_range(
 30        excerpt.range.clone(),
 31        excerpt_text.body.as_str(),
 32        ReferenceRegion::Nearby,
 33        snapshot,
 34    );
 35
 36    for ((_, range), text) in excerpt
 37        .parent_declarations
 38        .iter()
 39        .zip(excerpt_text.parent_signatures.iter())
 40    {
 41        references.extend(identifiers_in_range(
 42            range.clone(),
 43            text.as_str(),
 44            ReferenceRegion::Breadcrumb,
 45            snapshot,
 46        ));
 47    }
 48
 49    let mut identifier_to_references: HashMap<Identifier, Vec<Reference>> = HashMap::new();
 50    for reference in references {
 51        identifier_to_references
 52            .entry(reference.identifier.clone())
 53            .or_insert_with(Vec::new)
 54            .push(reference);
 55    }
 56    identifier_to_references
 57}
 58
 59/// Finds all nodes which have a "variable" match from the highlights query within the offset range.
 60pub fn identifiers_in_range(
 61    range: Range<usize>,
 62    range_text: &str,
 63    reference_region: ReferenceRegion,
 64    buffer: &BufferSnapshot,
 65) -> Vec<Reference> {
 66    let mut matches = buffer
 67        .syntax
 68        .matches(range.clone(), &buffer.text, |grammar| {
 69            grammar
 70                .highlights_config
 71                .as_ref()
 72                .map(|config| &config.query)
 73        });
 74
 75    let mut references = Vec::new();
 76    let mut last_added_range = None;
 77    while let Some(mat) = matches.peek() {
 78        let config = matches.grammars()[mat.grammar_index]
 79            .highlights_config
 80            .as_ref();
 81
 82        if let Some(config) = config {
 83            for capture in mat.captures {
 84                if config.identifier_capture_indices.contains(&capture.index) {
 85                    let node_range = capture.node.byte_range();
 86
 87                    // sometimes multiple highlight queries match - this deduplicates them
 88                    if Some(node_range.clone()) == last_added_range {
 89                        continue;
 90                    }
 91
 92                    if !range.contains_inclusive(&node_range) {
 93                        continue;
 94                    }
 95
 96                    let identifier_text =
 97                        &range_text[node_range.start - range.start..node_range.end - range.start];
 98
 99                    references.push(Reference {
100                        identifier: Identifier {
101                            name: identifier_text.into(),
102                            language_id: mat.language.id(),
103                        },
104                        range: node_range.clone(),
105                        region: reference_region,
106                    });
107                    last_added_range = Some(node_range);
108                }
109            }
110        }
111
112        matches.advance();
113    }
114    references
115}
116
117#[cfg(test)]
118mod test {
119    use gpui::{TestAppContext, prelude::*};
120    use indoc::indoc;
121    use language::{BufferSnapshot, Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
122
123    use crate::reference::{ReferenceRegion, identifiers_in_range};
124
125    #[gpui::test]
126    fn test_identifier_node_truncated(cx: &mut TestAppContext) {
127        let code = indoc! { r#"
128            fn main() {
129                add(1, 2);
130            }
131
132            fn add(a: i32, b: i32) -> i32 {
133                a + b
134            }
135        "# };
136        let buffer = create_buffer(code, cx);
137
138        let range = 0..35;
139        let references = identifiers_in_range(
140            range.clone(),
141            &code[range],
142            ReferenceRegion::Breadcrumb,
143            &buffer,
144        );
145        assert_eq!(references.len(), 2);
146        assert_eq!(references[0].identifier.name.as_ref(), "main");
147        assert_eq!(references[1].identifier.name.as_ref(), "add");
148    }
149
150    fn create_buffer(text: &str, cx: &mut TestAppContext) -> BufferSnapshot {
151        let buffer =
152            cx.new(|cx| language::Buffer::local(text, cx).with_language(rust_lang().into(), cx));
153        buffer.read_with(cx, |buffer, _| buffer.snapshot())
154    }
155
156    fn rust_lang() -> Language {
157        Language::new(
158            LanguageConfig {
159                name: "Rust".into(),
160                matcher: LanguageMatcher {
161                    path_suffixes: vec!["rs".to_string()],
162                    ..Default::default()
163                },
164                ..Default::default()
165            },
166            Some(tree_sitter_rust::LANGUAGE.into()),
167        )
168        .with_highlights_query(include_str!("../../languages/src/rust/highlights.scm"))
169        .unwrap()
170        .with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
171        .unwrap()
172    }
173}