completion_provider.rs

  1use std::ops::Range;
  2use std::path::Path;
  3use std::sync::Arc;
  4use std::sync::atomic::AtomicBool;
  5
  6use anyhow::Result;
  7use collections::HashMap;
  8use editor::display_map::CreaseId;
  9use editor::{CompletionProvider, Editor, ExcerptId};
 10use file_icons::FileIcons;
 11use gpui::{App, Entity, Task, WeakEntity};
 12use language::{Buffer, CodeLabel, HighlightId};
 13use lsp::CompletionContext;
 14use parking_lot::Mutex;
 15use project::{Completion, CompletionIntent, CompletionResponse, ProjectPath, WorktreeId};
 16use rope::Point;
 17use text::{Anchor, ToPoint};
 18use ui::prelude::*;
 19use workspace::Workspace;
 20
 21use crate::context_picker::MentionLink;
 22use crate::context_picker::file_context_picker::{extract_file_name_and_directory, search_files};
 23
 24#[derive(Default)]
 25pub struct MentionSet {
 26    paths_by_crease_id: HashMap<CreaseId, ProjectPath>,
 27}
 28
 29impl MentionSet {
 30    pub fn insert(&mut self, crease_id: CreaseId, path: ProjectPath) {
 31        self.paths_by_crease_id.insert(crease_id, path);
 32    }
 33
 34    pub fn path_for_crease_id(&self, crease_id: CreaseId) -> Option<ProjectPath> {
 35        self.paths_by_crease_id.get(&crease_id).cloned()
 36    }
 37
 38    pub fn drain(&mut self) -> impl Iterator<Item = CreaseId> {
 39        self.paths_by_crease_id.drain().map(|(id, _)| id)
 40    }
 41}
 42
 43pub struct ContextPickerCompletionProvider {
 44    workspace: WeakEntity<Workspace>,
 45    editor: WeakEntity<Editor>,
 46    mention_set: Arc<Mutex<MentionSet>>,
 47}
 48
 49impl ContextPickerCompletionProvider {
 50    pub fn new(
 51        mention_set: Arc<Mutex<MentionSet>>,
 52        workspace: WeakEntity<Workspace>,
 53        editor: WeakEntity<Editor>,
 54    ) -> Self {
 55        Self {
 56            mention_set,
 57            workspace,
 58            editor,
 59        }
 60    }
 61
 62    pub(crate) fn completion_for_path(
 63        project_path: ProjectPath,
 64        path_prefix: &str,
 65        is_recent: bool,
 66        is_directory: bool,
 67        excerpt_id: ExcerptId,
 68        source_range: Range<Anchor>,
 69        editor: Entity<Editor>,
 70        mention_set: Arc<Mutex<MentionSet>>,
 71        cx: &App,
 72    ) -> Completion {
 73        let (file_name, directory) =
 74            extract_file_name_and_directory(&project_path.path, path_prefix);
 75
 76        let label =
 77            build_code_label_for_full_path(&file_name, directory.as_ref().map(|s| s.as_ref()), cx);
 78        let full_path = if let Some(directory) = directory {
 79            format!("{}{}", directory, file_name)
 80        } else {
 81            file_name.to_string()
 82        };
 83
 84        let crease_icon_path = if is_directory {
 85            FileIcons::get_folder_icon(false, cx).unwrap_or_else(|| IconName::Folder.path().into())
 86        } else {
 87            FileIcons::get_icon(Path::new(&full_path), cx)
 88                .unwrap_or_else(|| IconName::File.path().into())
 89        };
 90        let completion_icon_path = if is_recent {
 91            IconName::HistoryRerun.path().into()
 92        } else {
 93            crease_icon_path.clone()
 94        };
 95
 96        let new_text = format!("{} ", MentionLink::for_file(&file_name, &full_path));
 97        let new_text_len = new_text.len();
 98        Completion {
 99            replace_range: source_range.clone(),
100            new_text,
101            label,
102            documentation: None,
103            source: project::CompletionSource::Custom,
104            icon_path: Some(completion_icon_path),
105            insert_text_mode: None,
106            confirm: Some(confirm_completion_callback(
107                crease_icon_path,
108                file_name,
109                project_path,
110                excerpt_id,
111                source_range.start,
112                new_text_len - 1,
113                editor,
114                mention_set,
115            )),
116        }
117    }
118}
119
120fn build_code_label_for_full_path(file_name: &str, directory: Option<&str>, cx: &App) -> CodeLabel {
121    let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
122    let mut label = CodeLabel::default();
123
124    label.push_str(&file_name, None);
125    label.push_str(" ", None);
126
127    if let Some(directory) = directory {
128        label.push_str(&directory, comment_id);
129    }
130
131    label.filter_range = 0..label.text().len();
132
133    label
134}
135
136impl CompletionProvider for ContextPickerCompletionProvider {
137    fn completions(
138        &self,
139        excerpt_id: ExcerptId,
140        buffer: &Entity<Buffer>,
141        buffer_position: Anchor,
142        _trigger: CompletionContext,
143        _window: &mut Window,
144        cx: &mut Context<Editor>,
145    ) -> Task<Result<Vec<CompletionResponse>>> {
146        let state = buffer.update(cx, |buffer, _cx| {
147            let position = buffer_position.to_point(buffer);
148            let line_start = Point::new(position.row, 0);
149            let offset_to_line = buffer.point_to_offset(line_start);
150            let mut lines = buffer.text_for_range(line_start..position).lines();
151            let line = lines.next()?;
152            MentionCompletion::try_parse(line, offset_to_line)
153        });
154        let Some(state) = state else {
155            return Task::ready(Ok(Vec::new()));
156        };
157
158        let Some(workspace) = self.workspace.upgrade() else {
159            return Task::ready(Ok(Vec::new()));
160        };
161
162        let snapshot = buffer.read(cx).snapshot();
163        let source_range = snapshot.anchor_before(state.source_range.start)
164            ..snapshot.anchor_after(state.source_range.end);
165
166        let editor = self.editor.clone();
167        let mention_set = self.mention_set.clone();
168        let MentionCompletion { argument, .. } = state;
169        let query = argument.unwrap_or_else(|| "".to_string());
170
171        let search_task = search_files(query.clone(), Arc::<AtomicBool>::default(), &workspace, cx);
172
173        cx.spawn(async move |_, cx| {
174            let matches = search_task.await;
175            let Some(editor) = editor.upgrade() else {
176                return Ok(Vec::new());
177            };
178
179            let completions = cx.update(|cx| {
180                matches
181                    .into_iter()
182                    .map(|mat| {
183                        let path_match = &mat.mat;
184                        let project_path = ProjectPath {
185                            worktree_id: WorktreeId::from_usize(path_match.worktree_id),
186                            path: path_match.path.clone(),
187                        };
188
189                        Self::completion_for_path(
190                            project_path,
191                            &path_match.path_prefix,
192                            mat.is_recent,
193                            path_match.is_dir,
194                            excerpt_id,
195                            source_range.clone(),
196                            editor.clone(),
197                            mention_set.clone(),
198                            cx,
199                        )
200                    })
201                    .collect()
202            })?;
203
204            Ok(vec![CompletionResponse {
205                completions,
206                // Since this does its own filtering (see `filter_completions()` returns false),
207                // there is no benefit to computing whether this set of completions is incomplete.
208                is_incomplete: true,
209            }])
210        })
211    }
212
213    fn is_completion_trigger(
214        &self,
215        buffer: &Entity<language::Buffer>,
216        position: language::Anchor,
217        _text: &str,
218        _trigger_in_words: bool,
219        _menu_is_open: bool,
220        cx: &mut Context<Editor>,
221    ) -> bool {
222        let buffer = buffer.read(cx);
223        let position = position.to_point(buffer);
224        let line_start = Point::new(position.row, 0);
225        let offset_to_line = buffer.point_to_offset(line_start);
226        let mut lines = buffer.text_for_range(line_start..position).lines();
227        if let Some(line) = lines.next() {
228            MentionCompletion::try_parse(line, offset_to_line)
229                .map(|completion| {
230                    completion.source_range.start <= offset_to_line + position.column as usize
231                        && completion.source_range.end >= offset_to_line + position.column as usize
232                })
233                .unwrap_or(false)
234        } else {
235            false
236        }
237    }
238
239    fn sort_completions(&self) -> bool {
240        false
241    }
242
243    fn filter_completions(&self) -> bool {
244        false
245    }
246}
247
248fn confirm_completion_callback(
249    crease_icon_path: SharedString,
250    crease_text: SharedString,
251    project_path: ProjectPath,
252    excerpt_id: ExcerptId,
253    start: Anchor,
254    content_len: usize,
255    editor: Entity<Editor>,
256    mention_set: Arc<Mutex<MentionSet>>,
257) -> Arc<dyn Fn(CompletionIntent, &mut Window, &mut App) -> bool + Send + Sync> {
258    Arc::new(move |_, window, cx| {
259        let crease_text = crease_text.clone();
260        let crease_icon_path = crease_icon_path.clone();
261        let editor = editor.clone();
262        let project_path = project_path.clone();
263        let mention_set = mention_set.clone();
264        window.defer(cx, move |window, cx| {
265            let crease_id = crate::context_picker::insert_crease_for_mention(
266                excerpt_id,
267                start,
268                content_len,
269                crease_text.clone(),
270                crease_icon_path,
271                editor.clone(),
272                window,
273                cx,
274            );
275            if let Some(crease_id) = crease_id {
276                mention_set.lock().insert(crease_id, project_path);
277            }
278        });
279        false
280    })
281}
282
283#[derive(Debug, Default, PartialEq)]
284struct MentionCompletion {
285    source_range: Range<usize>,
286    argument: Option<String>,
287}
288
289impl MentionCompletion {
290    fn try_parse(line: &str, offset_to_line: usize) -> Option<Self> {
291        let last_mention_start = line.rfind('@')?;
292        if last_mention_start >= line.len() {
293            return Some(Self::default());
294        }
295        if last_mention_start > 0
296            && line
297                .chars()
298                .nth(last_mention_start - 1)
299                .map_or(false, |c| !c.is_whitespace())
300        {
301            return None;
302        }
303
304        let rest_of_line = &line[last_mention_start + 1..];
305        let mut argument = None;
306
307        let mut parts = rest_of_line.split_whitespace();
308        let mut end = last_mention_start + 1;
309        if let Some(argument_text) = parts.next() {
310            end += argument_text.len();
311            argument = Some(argument_text.to_string());
312        }
313
314        Some(Self {
315            source_range: last_mention_start + offset_to_line..end + offset_to_line,
316            argument,
317        })
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use gpui::{EventEmitter, FocusHandle, Focusable, TestAppContext, VisualTestContext};
325    use project::{Project, ProjectPath};
326    use serde_json::json;
327    use settings::SettingsStore;
328    use std::{ops::Deref, rc::Rc};
329    use util::path;
330    use workspace::{AppState, Item};
331
332    #[test]
333    fn test_mention_completion_parse() {
334        assert_eq!(MentionCompletion::try_parse("Lorem Ipsum", 0), None);
335
336        assert_eq!(
337            MentionCompletion::try_parse("Lorem @", 0),
338            Some(MentionCompletion {
339                source_range: 6..7,
340                argument: None,
341            })
342        );
343
344        assert_eq!(
345            MentionCompletion::try_parse("Lorem @main", 0),
346            Some(MentionCompletion {
347                source_range: 6..11,
348                argument: Some("main".to_string()),
349            })
350        );
351
352        assert_eq!(MentionCompletion::try_parse("test@", 0), None);
353    }
354
355    struct AtMentionEditor(Entity<Editor>);
356
357    impl Item for AtMentionEditor {
358        type Event = ();
359
360        fn include_in_nav_history() -> bool {
361            false
362        }
363
364        fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
365            "Test".into()
366        }
367    }
368
369    impl EventEmitter<()> for AtMentionEditor {}
370
371    impl Focusable for AtMentionEditor {
372        fn focus_handle(&self, cx: &App) -> FocusHandle {
373            self.0.read(cx).focus_handle(cx).clone()
374        }
375    }
376
377    impl Render for AtMentionEditor {
378        fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
379            self.0.clone().into_any_element()
380        }
381    }
382
383    #[gpui::test]
384    async fn test_context_completion_provider(cx: &mut TestAppContext) {
385        init_test(cx);
386
387        let app_state = cx.update(AppState::test);
388
389        cx.update(|cx| {
390            language::init(cx);
391            editor::init(cx);
392            workspace::init(app_state.clone(), cx);
393            Project::init_settings(cx);
394        });
395
396        app_state
397            .fs
398            .as_fake()
399            .insert_tree(
400                path!("/dir"),
401                json!({
402                    "editor": "",
403                    "a": {
404                        "one.txt": "",
405                        "two.txt": "",
406                        "three.txt": "",
407                        "four.txt": ""
408                    },
409                    "b": {
410                        "five.txt": "",
411                        "six.txt": "",
412                        "seven.txt": "",
413                        "eight.txt": "",
414                    }
415                }),
416            )
417            .await;
418
419        let project = Project::test(app_state.fs.clone(), [path!("/dir").as_ref()], cx).await;
420        let window = cx.add_window(|window, cx| Workspace::test_new(project.clone(), window, cx));
421        let workspace = window.root(cx).unwrap();
422
423        let worktree = project.update(cx, |project, cx| {
424            let mut worktrees = project.worktrees(cx).collect::<Vec<_>>();
425            assert_eq!(worktrees.len(), 1);
426            worktrees.pop().unwrap()
427        });
428        let worktree_id = worktree.read_with(cx, |worktree, _| worktree.id());
429
430        let mut cx = VisualTestContext::from_window(*window.deref(), cx);
431
432        let paths = vec![
433            path!("a/one.txt"),
434            path!("a/two.txt"),
435            path!("a/three.txt"),
436            path!("a/four.txt"),
437            path!("b/five.txt"),
438            path!("b/six.txt"),
439            path!("b/seven.txt"),
440            path!("b/eight.txt"),
441        ];
442
443        let mut opened_editors = Vec::new();
444        for path in paths {
445            let buffer = workspace
446                .update_in(&mut cx, |workspace, window, cx| {
447                    workspace.open_path(
448                        ProjectPath {
449                            worktree_id,
450                            path: Path::new(path).into(),
451                        },
452                        None,
453                        false,
454                        window,
455                        cx,
456                    )
457                })
458                .await
459                .unwrap();
460            opened_editors.push(buffer);
461        }
462
463        let editor = workspace.update_in(&mut cx, |workspace, window, cx| {
464            let editor = cx.new(|cx| {
465                Editor::new(
466                    editor::EditorMode::full(),
467                    multi_buffer::MultiBuffer::build_simple("", cx),
468                    None,
469                    window,
470                    cx,
471                )
472            });
473            workspace.active_pane().update(cx, |pane, cx| {
474                pane.add_item(
475                    Box::new(cx.new(|_| AtMentionEditor(editor.clone()))),
476                    true,
477                    true,
478                    None,
479                    window,
480                    cx,
481                );
482            });
483            editor
484        });
485
486        let mention_set = Arc::new(Mutex::new(MentionSet::default()));
487
488        let editor_entity = editor.downgrade();
489        editor.update_in(&mut cx, |editor, window, cx| {
490            window.focus(&editor.focus_handle(cx));
491            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
492                mention_set.clone(),
493                workspace.downgrade(),
494                editor_entity,
495            ))));
496        });
497
498        cx.simulate_input("Lorem ");
499
500        editor.update(&mut cx, |editor, cx| {
501            assert_eq!(editor.text(cx), "Lorem ");
502            assert!(!editor.has_visible_completions_menu());
503        });
504
505        cx.simulate_input("@");
506
507        editor.update(&mut cx, |editor, cx| {
508            assert_eq!(editor.text(cx), "Lorem @");
509            assert!(editor.has_visible_completions_menu());
510            assert_eq!(
511                current_completion_labels(editor),
512                &[
513                    "eight.txt dir/b/",
514                    "seven.txt dir/b/",
515                    "six.txt dir/b/",
516                    "five.txt dir/b/",
517                    "four.txt dir/a/",
518                    "three.txt dir/a/",
519                    "two.txt dir/a/",
520                    "one.txt dir/a/",
521                    "dir ",
522                    "a dir/",
523                    "four.txt dir/a/",
524                    "one.txt dir/a/",
525                    "three.txt dir/a/",
526                    "two.txt dir/a/",
527                    "b dir/",
528                    "eight.txt dir/b/",
529                    "five.txt dir/b/",
530                    "seven.txt dir/b/",
531                    "six.txt dir/b/",
532                    "editor dir/"
533                ]
534            );
535        });
536
537        // Select and confirm "File"
538        editor.update_in(&mut cx, |editor, window, cx| {
539            assert!(editor.has_visible_completions_menu());
540            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
541            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
542            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
543            editor.context_menu_next(&editor::actions::ContextMenuNext, window, cx);
544            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
545        });
546
547        cx.run_until_parked();
548
549        editor.update(&mut cx, |editor, cx| {
550            assert_eq!(editor.text(cx), "Lorem [@four.txt](@file:dir/a/four.txt) ");
551        });
552    }
553
554    fn current_completion_labels(editor: &Editor) -> Vec<String> {
555        let completions = editor.current_completions().expect("Missing completions");
556        completions
557            .into_iter()
558            .map(|completion| completion.label.text.to_string())
559            .collect::<Vec<_>>()
560    }
561
562    pub(crate) fn init_test(cx: &mut TestAppContext) {
563        cx.update(|cx| {
564            let store = SettingsStore::test(cx);
565            cx.set_global(store);
566            theme::init(theme::LoadThemes::JustBase, cx);
567            client::init_settings(cx);
568            language::init(cx);
569            Project::init_settings(cx);
570            workspace::init_settings(cx);
571            editor::init_settings(cx);
572        });
573    }
574}