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 using tree-sitter,
 13/// making the assumption that all names are unique, and all variables' types are
 14/// 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                text_document_sync: Some(TextDocumentSyncCapability::Kind(
 38                    TextDocumentSyncKind::FULL,
 39                )),
 40                ..Default::default()
 41            },
 42            label_for_completion: None,
 43            initializer: Some(Box::new({
 44                move |server| {
 45                    server.handle_notification::<lsp::notification::DidOpenTextDocument, _>({
 46                        let index = index.clone();
 47                        move |params, _cx| {
 48                            index
 49                                .lock()
 50                                .open_buffer(params.text_document.uri, &params.text_document.text);
 51                        }
 52                    });
 53
 54                    server.handle_notification::<lsp::notification::DidCloseTextDocument, _>({
 55                        let index = index.clone();
 56                        let fs = fs.clone();
 57                        move |params, cx| {
 58                            let uri = params.text_document.uri;
 59                            let path = uri.to_file_path().ok();
 60                            index.lock().mark_buffer_closed(&uri);
 61
 62                            if let Some(path) = path {
 63                                let index = index.clone();
 64                                let fs = fs.clone();
 65                                cx.spawn(async move |_cx| {
 66                                    if let Ok(content) = fs.load(&path).await {
 67                                        index.lock().index_file(uri, &content);
 68                                    }
 69                                })
 70                                .detach();
 71                            }
 72                        }
 73                    });
 74
 75                    server.handle_notification::<lsp::notification::DidChangeWatchedFiles, _>({
 76                        let index = index.clone();
 77                        let fs = fs.clone();
 78                        move |params, cx| {
 79                            let index = index.clone();
 80                            let fs = fs.clone();
 81                            cx.spawn(async move |_cx| {
 82                                for event in params.changes {
 83                                    if index.lock().is_buffer_open(&event.uri) {
 84                                        continue;
 85                                    }
 86
 87                                    match event.typ {
 88                                        lsp::FileChangeType::DELETED => {
 89                                            index.lock().remove_definitions_for_file(&event.uri);
 90                                        }
 91                                        lsp::FileChangeType::CREATED
 92                                        | lsp::FileChangeType::CHANGED => {
 93                                            if let Some(path) = event.uri.to_file_path().ok() {
 94                                                if let Ok(content) = fs.load(&path).await {
 95                                                    index.lock().index_file(event.uri, &content);
 96                                                }
 97                                            }
 98                                        }
 99                                        _ => {}
100                                    }
101                                }
102                            })
103                            .detach();
104                        }
105                    });
106
107                    server.handle_notification::<lsp::notification::DidChangeTextDocument, _>({
108                        let index = index.clone();
109                        move |params, _cx| {
110                            if let Some(change) = params.content_changes.into_iter().last() {
111                                index
112                                    .lock()
113                                    .index_file(params.text_document.uri, &change.text);
114                            }
115                        }
116                    });
117
118                    server.handle_notification::<lsp::notification::DidChangeWorkspaceFolders, _>(
119                        {
120                            let index = index.clone();
121                            let fs = fs.clone();
122                            move |params, cx| {
123                                let index = index.clone();
124                                let fs = fs.clone();
125                                let files = fs.as_fake().files();
126                                cx.spawn(async move |_cx| {
127                                    for folder in params.event.added {
128                                        let Ok(path) = folder.uri.to_file_path() else {
129                                            continue;
130                                        };
131                                        for file in &files {
132                                            if let Some(uri) = Uri::from_file_path(&file).ok()
133                                                && file.starts_with(&path)
134                                                && let Ok(content) = fs.load(&file).await
135                                            {
136                                                index.lock().index_file(uri, &content);
137                                            }
138                                        }
139                                    }
140                                })
141                                .detach();
142                            }
143                        },
144                    );
145
146                    server.set_request_handler::<lsp::request::GotoDefinition, _, _>({
147                        let index = index.clone();
148                        move |params, _cx| {
149                            let result = index.lock().get_definitions(
150                                params.text_document_position_params.text_document.uri,
151                                params.text_document_position_params.position,
152                            );
153                            async move { Ok(result) }
154                        }
155                    });
156                }
157            })),
158        },
159    )
160}
161
162struct DefinitionIndex {
163    language: Arc<Language>,
164    definitions: HashMap<String, Vec<lsp::Location>>,
165    files: HashMap<Uri, FileEntry>,
166}
167
168#[derive(Debug)]
169struct FileEntry {
170    contents: String,
171    is_open_in_buffer: bool,
172}
173
174impl DefinitionIndex {
175    fn new(language: Arc<Language>) -> Self {
176        Self {
177            language,
178            definitions: HashMap::default(),
179            files: HashMap::default(),
180        }
181    }
182
183    fn remove_definitions_for_file(&mut self, uri: &Uri) {
184        self.definitions.retain(|_, locations| {
185            locations.retain(|loc| &loc.uri != uri);
186            !locations.is_empty()
187        });
188        self.files.remove(uri);
189    }
190
191    fn open_buffer(&mut self, uri: Uri, content: &str) {
192        self.index_file_inner(uri, content, true);
193    }
194
195    fn mark_buffer_closed(&mut self, uri: &Uri) {
196        if let Some(entry) = self.files.get_mut(uri) {
197            entry.is_open_in_buffer = false;
198        }
199    }
200
201    fn is_buffer_open(&self, uri: &Uri) -> bool {
202        self.files
203            .get(uri)
204            .map(|entry| entry.is_open_in_buffer)
205            .unwrap_or(false)
206    }
207
208    fn index_file(&mut self, uri: Uri, content: &str) {
209        self.index_file_inner(uri, content, false);
210    }
211
212    fn index_file_inner(&mut self, uri: Uri, content: &str, is_open_in_buffer: bool) -> Option<()> {
213        self.remove_definitions_for_file(&uri);
214        let grammar = self.language.grammar()?;
215        let outline_config = grammar.outline_config.as_ref()?;
216        let mut parser = Parser::new();
217        parser.set_language(&grammar.ts_language).ok()?;
218        let tree = parser.parse(content, None)?;
219        let declarations = extract_declarations_from_tree(&tree, content, outline_config);
220        for (name, byte_range) in declarations {
221            let range = byte_range_to_lsp_range(content, byte_range);
222            let location = lsp::Location {
223                uri: uri.clone(),
224                range,
225            };
226            self.definitions
227                .entry(name)
228                .or_insert_with(Vec::new)
229                .push(location);
230        }
231        self.files.insert(
232            uri,
233            FileEntry {
234                contents: content.to_string(),
235                is_open_in_buffer,
236            },
237        );
238
239        Some(())
240    }
241
242    fn get_definitions(
243        &mut self,
244        uri: Uri,
245        position: lsp::Position,
246    ) -> Option<lsp::GotoDefinitionResponse> {
247        let entry = self.files.get(&uri)?;
248        let name = word_at_position(&entry.contents, position)?;
249        let locations = self.definitions.get(name).cloned()?;
250        Some(lsp::GotoDefinitionResponse::Array(locations))
251    }
252}
253
254fn extract_declarations_from_tree(
255    tree: &Tree,
256    content: &str,
257    outline_config: &language::OutlineConfig,
258) -> Vec<(String, Range<usize>)> {
259    let mut cursor = QueryCursor::new();
260    let mut declarations = Vec::new();
261    let mut matches = cursor.matches(&outline_config.query, tree.root_node(), content.as_bytes());
262    while let Some(query_match) = matches.next() {
263        let mut name_range: Option<Range<usize>> = None;
264        let mut has_item_range = false;
265
266        for capture in query_match.captures {
267            let range = capture.node.byte_range();
268            if capture.index == outline_config.name_capture_ix {
269                name_range = Some(range);
270            } else if capture.index == outline_config.item_capture_ix {
271                has_item_range = true;
272            }
273        }
274
275        if let Some(name_range) = name_range
276            && has_item_range
277        {
278            let name = content[name_range.clone()].to_string();
279            if declarations.iter().any(|(n, _)| n == &name) {
280                continue;
281            }
282            declarations.push((name, name_range));
283        }
284    }
285    declarations
286}
287
288fn byte_range_to_lsp_range(content: &str, byte_range: Range<usize>) -> lsp::Range {
289    let start = byte_offset_to_position(content, byte_range.start);
290    let end = byte_offset_to_position(content, byte_range.end);
291    lsp::Range { start, end }
292}
293
294fn byte_offset_to_position(content: &str, offset: usize) -> lsp::Position {
295    let mut line = 0;
296    let mut character = 0;
297    let mut current_offset = 0;
298    for ch in content.chars() {
299        if current_offset >= offset {
300            break;
301        }
302        if ch == '\n' {
303            line += 1;
304            character = 0;
305        } else {
306            character += 1;
307        }
308        current_offset += ch.len_utf8();
309    }
310    lsp::Position { line, character }
311}
312
313fn word_at_position(content: &str, position: lsp::Position) -> Option<&str> {
314    let mut lines = content.lines();
315    let line = lines.nth(position.line as usize)?;
316    let column = position.character as usize;
317    if column > line.len() {
318        return None;
319    }
320    let start = line[..column]
321        .rfind(|c: char| !c.is_alphanumeric() && c != '_')
322        .map(|i| i + 1)
323        .unwrap_or(0);
324    let end = line[column..]
325        .find(|c: char| !c.is_alphanumeric() && c != '_')
326        .map(|i| i + column)
327        .unwrap_or(line.len());
328    Some(&line[start..end]).filter(|word| !word.is_empty())
329}