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_by_file: HashMap<Uri, 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_by_file: 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.type_annotations_by_file.remove(uri);
203        self.files.remove(uri);
204    }
205
206    fn open_buffer(&mut self, uri: Uri, content: &str) {
207        self.index_file_inner(uri, content, true);
208    }
209
210    fn mark_buffer_closed(&mut self, uri: &Uri) {
211        if let Some(entry) = self.files.get_mut(uri) {
212            entry.is_open_in_buffer = false;
213        }
214    }
215
216    fn is_buffer_open(&self, uri: &Uri) -> bool {
217        self.files
218            .get(uri)
219            .map(|entry| entry.is_open_in_buffer)
220            .unwrap_or(false)
221    }
222
223    fn index_file(&mut self, uri: Uri, content: &str) {
224        self.index_file_inner(uri, content, false);
225    }
226
227    fn index_file_inner(&mut self, uri: Uri, content: &str, is_open_in_buffer: bool) -> Option<()> {
228        self.remove_definitions_for_file(&uri);
229        let grammar = self.language.grammar()?;
230        let outline_config = grammar.outline_config.as_ref()?;
231        let mut parser = Parser::new();
232        parser.set_language(&grammar.ts_language).ok()?;
233        let tree = parser.parse(content, None)?;
234        let declarations = extract_declarations_from_tree(&tree, content, outline_config);
235        for (name, byte_range) in declarations {
236            let range = byte_range_to_lsp_range(content, byte_range);
237            let location = lsp::Location {
238                uri: uri.clone(),
239                range,
240            };
241            self.definitions
242                .entry(name)
243                .or_insert_with(Vec::new)
244                .push(location);
245        }
246
247        let type_annotations = extract_type_annotations(content)
248            .into_iter()
249            .collect::<HashMap<_, _>>();
250        self.type_annotations_by_file
251            .insert(uri.clone(), type_annotations);
252
253        self.files.insert(
254            uri,
255            FileEntry {
256                contents: content.to_string(),
257                is_open_in_buffer,
258            },
259        );
260
261        Some(())
262    }
263
264    fn get_definitions(
265        &mut self,
266        uri: Uri,
267        position: lsp::Position,
268    ) -> Option<lsp::GotoDefinitionResponse> {
269        let entry = self.files.get(&uri)?;
270        let name = word_at_position(&entry.contents, position)?;
271        let locations = self.definitions.get(name).cloned()?;
272        Some(lsp::GotoDefinitionResponse::Array(locations))
273    }
274
275    fn get_type_definitions(
276        &mut self,
277        uri: Uri,
278        position: lsp::Position,
279    ) -> Option<lsp::GotoDefinitionResponse> {
280        let entry = self.files.get(&uri)?;
281        let name = word_at_position(&entry.contents, position)?;
282
283        if let Some(type_name) = self
284            .type_annotations_by_file
285            .get(&uri)
286            .and_then(|annotations| annotations.get(name))
287        {
288            if let Some(locations) = self.definitions.get(type_name) {
289                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
290            }
291        }
292
293        // If the identifier itself is an uppercase name (a type), return its own definition.
294        // This mirrors real LSP behavior where GotoTypeDefinition on a type name
295        // resolves to that type's definition.
296        if name.starts_with(|c: char| c.is_uppercase()) {
297            if let Some(locations) = self.definitions.get(name) {
298                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
299            }
300        }
301
302        None
303    }
304}
305
306/// Extracts `identifier_name -> type_name` mappings from field declarations
307/// and function parameters. For example, `owner: Arc<Person>` produces
308/// `"owner" -> "Person"` by unwrapping common generic wrappers.
309fn extract_type_annotations(content: &str) -> Vec<(String, String)> {
310    let mut annotations = Vec::new();
311    for line in content.lines() {
312        let trimmed = line.trim();
313        if trimmed.starts_with("//")
314            || trimmed.starts_with("use ")
315            || trimmed.starts_with("pub use ")
316        {
317            continue;
318        }
319
320        let Some(colon_idx) = trimmed.find(':') else {
321            continue;
322        };
323
324        // The part before `:` should end with an identifier name.
325        let left = trimmed[..colon_idx].trim();
326        let Some(name) = left.split_whitespace().last() else {
327            continue;
328        };
329
330        if name.is_empty() || !name.chars().all(|c| c.is_alphanumeric() || c == '_') {
331            continue;
332        }
333
334        // Skip names that start uppercase — they're type names, not variables/fields.
335        if name.starts_with(|c: char| c.is_uppercase()) {
336            continue;
337        }
338
339        let right = trimmed[colon_idx + 1..].trim();
340        let type_name = extract_base_type_name(right);
341
342        if !type_name.is_empty() && type_name.starts_with(|c: char| c.is_uppercase()) {
343            annotations.push((name.to_string(), type_name));
344        }
345    }
346    annotations
347}
348
349/// Unwraps common generic wrappers (Arc, Box, Rc, Option, Vec) and trait
350/// object prefixes (dyn, impl) to find the concrete type name. For example:
351/// `Arc<Person>` → `"Person"`, `Box<dyn Trait>` → `"Trait"`.
352fn extract_base_type_name(type_str: &str) -> String {
353    let trimmed = type_str
354        .trim()
355        .trim_start_matches('&')
356        .trim_start_matches("mut ")
357        .trim_end_matches(',')
358        .trim_end_matches('{')
359        .trim_end_matches(')')
360        .trim()
361        .trim_start_matches("dyn ")
362        .trim_start_matches("impl ")
363        .trim();
364
365    if let Some(angle_start) = trimmed.find('<') {
366        let outer = &trimmed[..angle_start];
367        if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
368            let inner_end = trimmed.rfind('>').unwrap_or(trimmed.len());
369            let inner = &trimmed[angle_start + 1..inner_end];
370            return extract_base_type_name(inner);
371        }
372        return outer.to_string();
373    }
374
375    if let Some(call_start) = trimmed.find("::") {
376        let outer = &trimmed[..call_start];
377        if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
378            let rest = trimmed[call_start + 2..].trim_start();
379            if let Some(paren_start) = rest.find('(') {
380                let inner = &rest[paren_start + 1..];
381                let inner = inner.trim();
382                if !inner.is_empty() {
383                    return extract_base_type_name(inner);
384                }
385            }
386        }
387    }
388
389    trimmed
390        .split(|c: char| !c.is_alphanumeric() && c != '_')
391        .next()
392        .unwrap_or("")
393        .to_string()
394}
395
396fn extract_declarations_from_tree(
397    tree: &Tree,
398    content: &str,
399    outline_config: &language::OutlineConfig,
400) -> Vec<(String, Range<usize>)> {
401    let mut cursor = QueryCursor::new();
402    let mut declarations = Vec::new();
403    let mut matches = cursor.matches(&outline_config.query, tree.root_node(), content.as_bytes());
404    while let Some(query_match) = matches.next() {
405        let mut name_range: Option<Range<usize>> = None;
406        let mut has_item_range = false;
407
408        for capture in query_match.captures {
409            let range = capture.node.byte_range();
410            if capture.index == outline_config.name_capture_ix {
411                name_range = Some(range);
412            } else if capture.index == outline_config.item_capture_ix {
413                has_item_range = true;
414            }
415        }
416
417        if let Some(name_range) = name_range
418            && has_item_range
419        {
420            let name = content[name_range.clone()].to_string();
421            if declarations.iter().any(|(n, _)| n == &name) {
422                continue;
423            }
424            declarations.push((name, name_range));
425        }
426    }
427    declarations
428}
429
430fn byte_range_to_lsp_range(content: &str, byte_range: Range<usize>) -> lsp::Range {
431    let start = byte_offset_to_position(content, byte_range.start);
432    let end = byte_offset_to_position(content, byte_range.end);
433    lsp::Range { start, end }
434}
435
436fn byte_offset_to_position(content: &str, offset: usize) -> lsp::Position {
437    let mut line = 0;
438    let mut character = 0;
439    let mut current_offset = 0;
440    for ch in content.chars() {
441        if current_offset >= offset {
442            break;
443        }
444        if ch == '\n' {
445            line += 1;
446            character = 0;
447        } else {
448            character += 1;
449        }
450        current_offset += ch.len_utf8();
451    }
452    lsp::Position { line, character }
453}
454
455fn word_at_position(content: &str, position: lsp::Position) -> Option<&str> {
456    let mut lines = content.lines();
457    let line = lines.nth(position.line as usize)?;
458    let column = position.character as usize;
459    if column > line.len() {
460        return None;
461    }
462    let start = line[..column]
463        .rfind(|c: char| !c.is_alphanumeric() && c != '_')
464        .map(|i| i + 1)
465        .unwrap_or(0);
466    let end = line[column..]
467        .find(|c: char| !c.is_alphanumeric() && c != '_')
468        .map(|i| i + column)
469        .unwrap_or(line.len());
470    Some(&line[start..end]).filter(|word| !word.is_empty())
471}