fake_definition_lsp.rs

  1use collections::HashMap;
  2use futures::channel::mpsc::UnboundedReceiver;
  3use language::{Language, LanguageRegistry};
  4use lsp::{
  5    FakeLanguageServer, LanguageServerBinary, TextDocumentSyncCapability, TextDocumentSyncKind, Uri,
  6};
  7use parking_lot::Mutex;
  8use project::Fs;
  9use std::{ops::Range, path::PathBuf, sync::Arc};
 10use tree_sitter::{Parser, QueryCursor, StreamingIterator, Tree};
 11
 12/// Registers a fake language server that implements go-to-definition and
 13/// go-to-type-definition using tree-sitter, making the assumption that all
 14/// names are unique, and all variables' types are explicitly declared.
 15pub fn register_fake_definition_server(
 16    language_registry: &Arc<LanguageRegistry>,
 17    language: Arc<Language>,
 18    fs: Arc<dyn Fs>,
 19) -> UnboundedReceiver<FakeLanguageServer> {
 20    let index = Arc::new(Mutex::new(DefinitionIndex::new(language.clone())));
 21
 22    language_registry.register_fake_lsp(
 23        language.name(),
 24        language::FakeLspAdapter {
 25            name: "fake-definition-lsp",
 26            initialization_options: None,
 27            prettier_plugins: Vec::new(),
 28            disk_based_diagnostics_progress_token: None,
 29            disk_based_diagnostics_sources: Vec::new(),
 30            language_server_binary: LanguageServerBinary {
 31                path: PathBuf::from("fake-definition-lsp"),
 32                arguments: Vec::new(),
 33                env: None,
 34            },
 35            capabilities: lsp::ServerCapabilities {
 36                definition_provider: Some(lsp::OneOf::Left(true)),
 37                type_definition_provider: Some(lsp::TypeDefinitionProviderCapability::Simple(true)),
 38                text_document_sync: Some(TextDocumentSyncCapability::Kind(
 39                    TextDocumentSyncKind::FULL,
 40                )),
 41                ..Default::default()
 42            },
 43            label_for_completion: None,
 44            initializer: Some(Box::new({
 45                move |server| {
 46                    server.handle_notification::<lsp::notification::DidOpenTextDocument, _>({
 47                        let index = index.clone();
 48                        move |params, _cx| {
 49                            index
 50                                .lock()
 51                                .open_buffer(params.text_document.uri, &params.text_document.text);
 52                        }
 53                    });
 54
 55                    server.handle_notification::<lsp::notification::DidCloseTextDocument, _>({
 56                        let index = index.clone();
 57                        let fs = fs.clone();
 58                        move |params, cx| {
 59                            let uri = params.text_document.uri;
 60                            let path = uri.to_file_path().ok();
 61                            index.lock().mark_buffer_closed(&uri);
 62
 63                            if let Some(path) = path {
 64                                let index = index.clone();
 65                                let fs = fs.clone();
 66                                cx.spawn(async move |_cx| {
 67                                    if let Ok(content) = fs.load(&path).await {
 68                                        index.lock().index_file(uri, &content);
 69                                    }
 70                                })
 71                                .detach();
 72                            }
 73                        }
 74                    });
 75
 76                    server.handle_notification::<lsp::notification::DidChangeWatchedFiles, _>({
 77                        let index = index.clone();
 78                        let fs = fs.clone();
 79                        move |params, cx| {
 80                            let index = index.clone();
 81                            let fs = fs.clone();
 82                            cx.spawn(async move |_cx| {
 83                                for event in params.changes {
 84                                    if index.lock().is_buffer_open(&event.uri) {
 85                                        continue;
 86                                    }
 87
 88                                    match event.typ {
 89                                        lsp::FileChangeType::DELETED => {
 90                                            index.lock().remove_definitions_for_file(&event.uri);
 91                                        }
 92                                        lsp::FileChangeType::CREATED
 93                                        | lsp::FileChangeType::CHANGED => {
 94                                            if let Some(path) = event.uri.to_file_path().ok() {
 95                                                if let Ok(content) = fs.load(&path).await {
 96                                                    index.lock().index_file(event.uri, &content);
 97                                                }
 98                                            }
 99                                        }
100                                        _ => {}
101                                    }
102                                }
103                            })
104                            .detach();
105                        }
106                    });
107
108                    server.handle_notification::<lsp::notification::DidChangeTextDocument, _>({
109                        let index = index.clone();
110                        move |params, _cx| {
111                            if let Some(change) = params.content_changes.into_iter().last() {
112                                index
113                                    .lock()
114                                    .index_file(params.text_document.uri, &change.text);
115                            }
116                        }
117                    });
118
119                    server.handle_notification::<lsp::notification::DidChangeWorkspaceFolders, _>(
120                        {
121                            let index = index.clone();
122                            let fs = fs.clone();
123                            move |params, cx| {
124                                let index = index.clone();
125                                let fs = fs.clone();
126                                let files = fs.as_fake().files();
127                                cx.spawn(async move |_cx| {
128                                    for folder in params.event.added {
129                                        let Ok(path) = folder.uri.to_file_path() else {
130                                            continue;
131                                        };
132                                        for file in &files {
133                                            if let Some(uri) = Uri::from_file_path(&file).ok()
134                                                && file.starts_with(&path)
135                                                && let Ok(content) = fs.load(&file).await
136                                            {
137                                                index.lock().index_file(uri, &content);
138                                            }
139                                        }
140                                    }
141                                })
142                                .detach();
143                            }
144                        },
145                    );
146
147                    server.set_request_handler::<lsp::request::GotoDefinition, _, _>({
148                        let index = index.clone();
149                        move |params, _cx| {
150                            let result = index.lock().get_definitions(
151                                params.text_document_position_params.text_document.uri,
152                                params.text_document_position_params.position,
153                            );
154                            async move { Ok(result) }
155                        }
156                    });
157
158                    server.set_request_handler::<lsp::request::GotoTypeDefinition, _, _>({
159                        let index = index.clone();
160                        move |params, _cx| {
161                            let result = index.lock().get_type_definitions(
162                                params.text_document_position_params.text_document.uri,
163                                params.text_document_position_params.position,
164                            );
165                            async move { Ok(result) }
166                        }
167                    });
168                }
169            })),
170        },
171    )
172}
173
174struct DefinitionIndex {
175    language: Arc<Language>,
176    definitions: HashMap<String, Vec<lsp::Location>>,
177    type_annotations: HashMap<String, String>,
178    files: HashMap<Uri, FileEntry>,
179}
180
181#[derive(Debug)]
182struct FileEntry {
183    contents: String,
184    is_open_in_buffer: bool,
185}
186
187impl DefinitionIndex {
188    fn new(language: Arc<Language>) -> Self {
189        Self {
190            language,
191            definitions: HashMap::default(),
192            type_annotations: HashMap::default(),
193            files: HashMap::default(),
194        }
195    }
196
197    fn remove_definitions_for_file(&mut self, uri: &Uri) {
198        self.definitions.retain(|_, locations| {
199            locations.retain(|loc| &loc.uri != uri);
200            !locations.is_empty()
201        });
202        self.files.remove(uri);
203    }
204
205    fn open_buffer(&mut self, uri: Uri, content: &str) {
206        self.index_file_inner(uri, content, true);
207    }
208
209    fn mark_buffer_closed(&mut self, uri: &Uri) {
210        if let Some(entry) = self.files.get_mut(uri) {
211            entry.is_open_in_buffer = false;
212        }
213    }
214
215    fn is_buffer_open(&self, uri: &Uri) -> bool {
216        self.files
217            .get(uri)
218            .map(|entry| entry.is_open_in_buffer)
219            .unwrap_or(false)
220    }
221
222    fn index_file(&mut self, uri: Uri, content: &str) {
223        self.index_file_inner(uri, content, false);
224    }
225
226    fn index_file_inner(&mut self, uri: Uri, content: &str, is_open_in_buffer: bool) -> Option<()> {
227        self.remove_definitions_for_file(&uri);
228        let grammar = self.language.grammar()?;
229        let outline_config = grammar.outline_config.as_ref()?;
230        let mut parser = Parser::new();
231        parser.set_language(&grammar.ts_language).ok()?;
232        let tree = parser.parse(content, None)?;
233        let declarations = extract_declarations_from_tree(&tree, content, outline_config);
234        for (name, byte_range) in declarations {
235            let range = byte_range_to_lsp_range(content, byte_range);
236            let location = lsp::Location {
237                uri: uri.clone(),
238                range,
239            };
240            self.definitions
241                .entry(name)
242                .or_insert_with(Vec::new)
243                .push(location);
244        }
245
246        for (identifier_name, type_name) in extract_type_annotations(content) {
247            self.type_annotations
248                .entry(identifier_name)
249                .or_insert(type_name);
250        }
251
252        self.files.insert(
253            uri,
254            FileEntry {
255                contents: content.to_string(),
256                is_open_in_buffer,
257            },
258        );
259
260        Some(())
261    }
262
263    fn get_definitions(
264        &mut self,
265        uri: Uri,
266        position: lsp::Position,
267    ) -> Option<lsp::GotoDefinitionResponse> {
268        let entry = self.files.get(&uri)?;
269        let name = word_at_position(&entry.contents, position)?;
270        let locations = self.definitions.get(name).cloned()?;
271        Some(lsp::GotoDefinitionResponse::Array(locations))
272    }
273
274    fn get_type_definitions(
275        &mut self,
276        uri: Uri,
277        position: lsp::Position,
278    ) -> Option<lsp::GotoDefinitionResponse> {
279        let entry = self.files.get(&uri)?;
280        let name = word_at_position(&entry.contents, position)?;
281
282        if let Some(type_name) = self.type_annotations.get(name) {
283            if let Some(locations) = self.definitions.get(type_name) {
284                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
285            }
286        }
287
288        // If the identifier itself is an uppercase name (a type), return its own definition.
289        // This mirrors real LSP behavior where GotoTypeDefinition on a type name
290        // resolves to that type's definition.
291        if name.starts_with(|c: char| c.is_uppercase()) {
292            if let Some(locations) = self.definitions.get(name) {
293                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
294            }
295        }
296
297        None
298    }
299}
300
301/// Extracts `identifier_name -> type_name` mappings from field declarations
302/// and function parameters. For example, `owner: Arc<Person>` produces
303/// `"owner" -> "Person"` by unwrapping common generic wrappers.
304fn extract_type_annotations(content: &str) -> Vec<(String, String)> {
305    let mut annotations = Vec::new();
306    for line in content.lines() {
307        let trimmed = line.trim();
308        if trimmed.starts_with("//")
309            || trimmed.starts_with("use ")
310            || trimmed.starts_with("pub use ")
311        {
312            continue;
313        }
314
315        let Some(colon_idx) = trimmed.find(':') else {
316            continue;
317        };
318
319        // The part before `:` should end with an identifier name.
320        let left = trimmed[..colon_idx].trim();
321        let Some(name) = left.split_whitespace().last() else {
322            continue;
323        };
324
325        if name.is_empty() || !name.chars().all(|c| c.is_alphanumeric() || c == '_') {
326            continue;
327        }
328
329        // Skip names that start uppercase — they're type names, not variables/fields.
330        if name.starts_with(|c: char| c.is_uppercase()) {
331            continue;
332        }
333
334        let right = trimmed[colon_idx + 1..].trim();
335        let type_name = extract_base_type_name(right);
336
337        if !type_name.is_empty() && type_name.starts_with(|c: char| c.is_uppercase()) {
338            annotations.push((name.to_string(), type_name));
339        }
340    }
341    annotations
342}
343
344/// Unwraps common generic wrappers (Arc, Box, Rc, Option, Vec) and trait
345/// object prefixes (dyn, impl) to find the concrete type name. For example:
346/// `Arc<Person>` → `"Person"`, `Box<dyn Trait>` → `"Trait"`.
347fn extract_base_type_name(type_str: &str) -> String {
348    let trimmed = type_str
349        .trim()
350        .trim_start_matches('&')
351        .trim_start_matches("mut ")
352        .trim_end_matches(',')
353        .trim_end_matches('{')
354        .trim_end_matches(')')
355        .trim()
356        .trim_start_matches("dyn ")
357        .trim_start_matches("impl ")
358        .trim();
359
360    if let Some(angle_start) = trimmed.find('<') {
361        let outer = &trimmed[..angle_start];
362        if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
363            let inner_end = trimmed.rfind('>').unwrap_or(trimmed.len());
364            let inner = &trimmed[angle_start + 1..inner_end];
365            return extract_base_type_name(inner);
366        }
367        return outer.to_string();
368    }
369
370    trimmed
371        .split(|c: char| !c.is_alphanumeric() && c != '_')
372        .next()
373        .unwrap_or("")
374        .to_string()
375}
376
377fn extract_declarations_from_tree(
378    tree: &Tree,
379    content: &str,
380    outline_config: &language::OutlineConfig,
381) -> Vec<(String, Range<usize>)> {
382    let mut cursor = QueryCursor::new();
383    let mut declarations = Vec::new();
384    let mut matches = cursor.matches(&outline_config.query, tree.root_node(), content.as_bytes());
385    while let Some(query_match) = matches.next() {
386        let mut name_range: Option<Range<usize>> = None;
387        let mut has_item_range = false;
388
389        for capture in query_match.captures {
390            let range = capture.node.byte_range();
391            if capture.index == outline_config.name_capture_ix {
392                name_range = Some(range);
393            } else if capture.index == outline_config.item_capture_ix {
394                has_item_range = true;
395            }
396        }
397
398        if let Some(name_range) = name_range
399            && has_item_range
400        {
401            let name = content[name_range.clone()].to_string();
402            if declarations.iter().any(|(n, _)| n == &name) {
403                continue;
404            }
405            declarations.push((name, name_range));
406        }
407    }
408    declarations
409}
410
411fn byte_range_to_lsp_range(content: &str, byte_range: Range<usize>) -> lsp::Range {
412    let start = byte_offset_to_position(content, byte_range.start);
413    let end = byte_offset_to_position(content, byte_range.end);
414    lsp::Range { start, end }
415}
416
417fn byte_offset_to_position(content: &str, offset: usize) -> lsp::Position {
418    let mut line = 0;
419    let mut character = 0;
420    let mut current_offset = 0;
421    for ch in content.chars() {
422        if current_offset >= offset {
423            break;
424        }
425        if ch == '\n' {
426            line += 1;
427            character = 0;
428        } else {
429            character += 1;
430        }
431        current_offset += ch.len_utf8();
432    }
433    lsp::Position { line, character }
434}
435
436fn word_at_position(content: &str, position: lsp::Position) -> Option<&str> {
437    let mut lines = content.lines();
438    let line = lines.nth(position.line as usize)?;
439    let column = position.character as usize;
440    if column > line.len() {
441        return None;
442    }
443    let start = line[..column]
444        .rfind(|c: char| !c.is_alphanumeric() && c != '_')
445        .map(|i| i + 1)
446        .unwrap_or(0);
447    let end = line[column..]
448        .find(|c: char| !c.is_alphanumeric() && c != '_')
449        .map(|i| i + column)
450        .unwrap_or(line.len());
451    Some(&line[start..end]).filter(|word| !word.is_empty())
452}