Skip to main content

emmylua_code_analysis/db_index/reference/
mod.rs

1mod file_reference;
2mod string_reference;
3
4use std::collections::{HashMap, HashSet};
5
6use emmylua_parser::LuaSyntaxId;
7pub use file_reference::{DeclReference, DeclReferenceCell, FileReference};
8use rowan::TextRange;
9use smol_str::SmolStr;
10use string_reference::StringReference;
11
12use super::{LuaDeclId, LuaMemberKey, LuaTypeDeclId, traits::LuaIndex};
13use crate::{FileId, InFiled};
14
15#[derive(Debug)]
16pub struct LuaReferenceIndex {
17    file_references: HashMap<FileId, FileReference>,
18    index_reference: HashMap<LuaMemberKey, HashMap<FileId, HashSet<LuaSyntaxId>>>,
19    global_references: HashMap<SmolStr, HashMap<FileId, HashSet<LuaSyntaxId>>>,
20    string_references: HashMap<FileId, StringReference>,
21    type_references: HashMap<FileId, HashMap<LuaTypeDeclId, HashSet<TextRange>>>,
22}
23
24impl Default for LuaReferenceIndex {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl LuaReferenceIndex {
31    pub fn new() -> Self {
32        Self {
33            file_references: HashMap::new(),
34            index_reference: HashMap::new(),
35            global_references: HashMap::new(),
36            string_references: HashMap::new(),
37            type_references: HashMap::new(),
38        }
39    }
40
41    pub fn add_decl_reference(
42        &mut self,
43        decl_id: LuaDeclId,
44        file_id: FileId,
45        range: TextRange,
46        is_write: bool,
47    ) {
48        self.file_references
49            .entry(file_id)
50            .or_default()
51            .add_decl_reference(decl_id, range, is_write);
52    }
53
54    pub fn add_global_reference(&mut self, name: &str, file_id: FileId, syntax_id: LuaSyntaxId) {
55        let key = SmolStr::new(name);
56        self.global_references
57            .entry(key)
58            .or_default()
59            .entry(file_id)
60            .or_default()
61            .insert(syntax_id);
62    }
63
64    pub fn add_index_reference(
65        &mut self,
66        key: LuaMemberKey,
67        file_id: FileId,
68        syntax_id: LuaSyntaxId,
69    ) {
70        self.index_reference
71            .entry(key)
72            .or_default()
73            .entry(file_id)
74            .or_default()
75            .insert(syntax_id);
76    }
77
78    pub fn add_string_reference(&mut self, file_id: FileId, string: &str, range: TextRange) {
79        self.string_references
80            .entry(file_id)
81            .or_insert_with(StringReference::new)
82            .add_string_reference(string, range);
83    }
84
85    pub fn add_type_reference(
86        &mut self,
87        file_id: FileId,
88        type_decl_id: LuaTypeDeclId,
89        range: TextRange,
90    ) {
91        self.type_references
92            .entry(file_id)
93            .or_default()
94            .entry(type_decl_id)
95            .or_default()
96            .insert(range);
97    }
98
99    pub fn get_local_reference(&self, file_id: &FileId) -> Option<&FileReference> {
100        self.file_references.get(file_id)
101    }
102
103    pub fn create_local_reference(&mut self, file_id: FileId) {
104        self.file_references.entry(file_id).or_default();
105    }
106
107    pub fn get_decl_references(
108        &self,
109        file_id: &FileId,
110        decl_id: &LuaDeclId,
111    ) -> Option<&DeclReference> {
112        self.file_references
113            .get(file_id)?
114            .get_decl_references(decl_id)
115    }
116
117    pub fn get_var_reference_decl(&self, file_id: &FileId, range: TextRange) -> Option<LuaDeclId> {
118        self.file_references.get(file_id)?.get_decl_id(&range)
119    }
120
121    pub fn get_decl_references_map(
122        &self,
123        file_id: &FileId,
124    ) -> Option<&HashMap<LuaDeclId, DeclReference>> {
125        self.file_references
126            .get(file_id)
127            .map(|file_reference| file_reference.get_decl_references_map())
128    }
129
130    pub fn get_global_file_references(
131        &self,
132        name: &str,
133        file_id: FileId,
134    ) -> Option<Vec<LuaSyntaxId>> {
135        let results = self
136            .global_references
137            .get(name)?
138            .iter()
139            .filter_map(|(source_file_id, syntax_ids)| {
140                if file_id == *source_file_id {
141                    Some(syntax_ids.iter())
142                } else {
143                    None
144                }
145            })
146            .flatten()
147            .copied()
148            .collect();
149
150        Some(results)
151    }
152
153    pub fn get_global_references(&self, name: &str) -> Option<Vec<InFiled<LuaSyntaxId>>> {
154        let results = self
155            .global_references
156            .get(name)?
157            .iter()
158            .flat_map(|(file_id, syntax_ids)| {
159                syntax_ids
160                    .iter()
161                    .map(|syntax_id| InFiled::new(*file_id, *syntax_id))
162            })
163            .collect();
164
165        Some(results)
166    }
167
168    pub fn get_index_references(&self, key: &LuaMemberKey) -> Option<Vec<InFiled<LuaSyntaxId>>> {
169        let results = self
170            .index_reference
171            .get(key)?
172            .iter()
173            .flat_map(|(file_id, syntax_ids)| {
174                syntax_ids
175                    .iter()
176                    .map(|syntax_id| InFiled::new(*file_id, *syntax_id))
177            })
178            .collect();
179
180        Some(results)
181    }
182
183    pub fn get_string_references(&self, string_value: &str) -> Vec<InFiled<TextRange>> {
184        self.string_references
185            .iter()
186            .flat_map(|(file_id, string_reference)| {
187                string_reference
188                    .get_string_references(string_value)
189                    .into_iter()
190                    .map(|range| InFiled::new(*file_id, range))
191            })
192            .collect()
193    }
194
195    pub fn get_type_references(
196        &self,
197        type_decl_id: &LuaTypeDeclId,
198    ) -> Option<Vec<InFiled<TextRange>>> {
199        let results = self
200            .type_references
201            .iter()
202            .flat_map(|(file_id, type_references)| {
203                type_references
204                    .get(type_decl_id)
205                    .into_iter()
206                    .flatten()
207                    .map(|range| InFiled::new(*file_id, *range))
208            })
209            .collect();
210
211        Some(results)
212    }
213}
214
215impl LuaIndex for LuaReferenceIndex {
216    fn remove(&mut self, file_id: FileId) {
217        self.file_references.remove(&file_id);
218        self.string_references.remove(&file_id);
219        self.type_references.remove(&file_id);
220        let mut to_be_remove = Vec::new();
221        for (key, references) in self.index_reference.iter_mut() {
222            references.remove(&file_id);
223            if references.is_empty() {
224                to_be_remove.push(key.clone());
225            }
226        }
227
228        for key in to_be_remove {
229            self.index_reference.remove(&key);
230        }
231
232        let mut to_be_remove = Vec::new();
233        for (key, references) in self.global_references.iter_mut() {
234            references.remove(&file_id);
235            if references.is_empty() {
236                to_be_remove.push(key.clone());
237            }
238        }
239
240        for key in to_be_remove {
241            self.global_references.remove(&key);
242        }
243    }
244
245    fn clear(&mut self) {
246        self.file_references.clear();
247        self.string_references.clear();
248        self.index_reference.clear();
249        self.global_references.clear();
250    }
251}