directory_context_picker.rs

  1use std::path::Path;
  2use std::sync::atomic::AtomicBool;
  3use std::sync::Arc;
  4
  5use anyhow::anyhow;
  6use fuzzy::PathMatch;
  7use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView};
  8use picker::{Picker, PickerDelegate};
  9use project::{PathMatchCandidateSet, ProjectPath, Worktree, WorktreeId};
 10use ui::{prelude::*, ListItem};
 11use util::ResultExt as _;
 12use workspace::Workspace;
 13
 14use crate::context::ContextKind;
 15use crate::context_picker::file_context_picker::codeblock_fence_for_path;
 16use crate::context_picker::{ConfirmBehavior, ContextPicker};
 17use crate::context_store::ContextStore;
 18
 19pub struct DirectoryContextPicker {
 20    picker: View<Picker<DirectoryContextPickerDelegate>>,
 21}
 22
 23impl DirectoryContextPicker {
 24    pub fn new(
 25        context_picker: WeakView<ContextPicker>,
 26        workspace: WeakView<Workspace>,
 27        context_store: WeakModel<ContextStore>,
 28        confirm_behavior: ConfirmBehavior,
 29        cx: &mut ViewContext<Self>,
 30    ) -> Self {
 31        let delegate = DirectoryContextPickerDelegate::new(
 32            context_picker,
 33            workspace,
 34            context_store,
 35            confirm_behavior,
 36        );
 37        let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
 38
 39        Self { picker }
 40    }
 41}
 42
 43impl FocusableView for DirectoryContextPicker {
 44    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 45        self.picker.focus_handle(cx)
 46    }
 47}
 48
 49impl Render for DirectoryContextPicker {
 50    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
 51        self.picker.clone()
 52    }
 53}
 54
 55pub struct DirectoryContextPickerDelegate {
 56    context_picker: WeakView<ContextPicker>,
 57    workspace: WeakView<Workspace>,
 58    context_store: WeakModel<ContextStore>,
 59    confirm_behavior: ConfirmBehavior,
 60    matches: Vec<PathMatch>,
 61    selected_index: usize,
 62}
 63
 64impl DirectoryContextPickerDelegate {
 65    pub fn new(
 66        context_picker: WeakView<ContextPicker>,
 67        workspace: WeakView<Workspace>,
 68        context_store: WeakModel<ContextStore>,
 69        confirm_behavior: ConfirmBehavior,
 70    ) -> Self {
 71        Self {
 72            context_picker,
 73            workspace,
 74            context_store,
 75            confirm_behavior,
 76            matches: Vec::new(),
 77            selected_index: 0,
 78        }
 79    }
 80
 81    fn search(
 82        &mut self,
 83        query: String,
 84        cancellation_flag: Arc<AtomicBool>,
 85        workspace: &View<Workspace>,
 86        cx: &mut ViewContext<Picker<Self>>,
 87    ) -> Task<Vec<PathMatch>> {
 88        if query.is_empty() {
 89            let workspace = workspace.read(cx);
 90            let project = workspace.project().read(cx);
 91            let directory_matches = project.worktrees(cx).flat_map(|worktree| {
 92                let worktree = worktree.read(cx);
 93                let path_prefix: Arc<str> = worktree.root_name().into();
 94                worktree.directories(false, 0).map(move |entry| PathMatch {
 95                    score: 0.,
 96                    positions: Vec::new(),
 97                    worktree_id: worktree.id().to_usize(),
 98                    path: entry.path.clone(),
 99                    path_prefix: path_prefix.clone(),
100                    distance_to_relative_ancestor: 0,
101                    is_dir: true,
102                })
103            });
104
105            Task::ready(directory_matches.collect())
106        } else {
107            let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
108            let candidate_sets = worktrees
109                .into_iter()
110                .map(|worktree| {
111                    let worktree = worktree.read(cx);
112
113                    PathMatchCandidateSet {
114                        snapshot: worktree.snapshot(),
115                        include_ignored: worktree
116                            .root_entry()
117                            .map_or(false, |entry| entry.is_ignored),
118                        include_root_name: true,
119                        candidates: project::Candidates::Directories,
120                    }
121                })
122                .collect::<Vec<_>>();
123
124            let executor = cx.background_executor().clone();
125            cx.foreground_executor().spawn(async move {
126                fuzzy::match_path_sets(
127                    candidate_sets.as_slice(),
128                    query.as_str(),
129                    None,
130                    false,
131                    100,
132                    &cancellation_flag,
133                    executor,
134                )
135                .await
136            })
137        }
138    }
139}
140
141impl PickerDelegate for DirectoryContextPickerDelegate {
142    type ListItem = ListItem;
143
144    fn match_count(&self) -> usize {
145        self.matches.len()
146    }
147
148    fn selected_index(&self) -> usize {
149        self.selected_index
150    }
151
152    fn set_selected_index(&mut self, ix: usize, _cx: &mut ViewContext<Picker<Self>>) {
153        self.selected_index = ix;
154    }
155
156    fn placeholder_text(&self, _cx: &mut WindowContext) -> Arc<str> {
157        "Search folders…".into()
158    }
159
160    fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
161        let Some(workspace) = self.workspace.upgrade() else {
162            return Task::ready(());
163        };
164
165        let search_task = self.search(query, Arc::<AtomicBool>::default(), &workspace, cx);
166
167        cx.spawn(|this, mut cx| async move {
168            let mut paths = search_task.await;
169            let empty_path = Path::new("");
170            paths.retain(|path_match| path_match.path.as_ref() != empty_path);
171
172            this.update(&mut cx, |this, _cx| {
173                this.delegate.matches = paths;
174            })
175            .log_err();
176        })
177    }
178
179    fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
180        let Some(mat) = self.matches.get(self.selected_index) else {
181            return;
182        };
183
184        let workspace = self.workspace.clone();
185        let Some(project) = workspace
186            .upgrade()
187            .map(|workspace| workspace.read(cx).project().clone())
188        else {
189            return;
190        };
191        let path = mat.path.clone();
192        let worktree_id = WorktreeId::from_usize(mat.worktree_id);
193        let confirm_behavior = self.confirm_behavior;
194        cx.spawn(|this, mut cx| async move {
195            let worktree = project.update(&mut cx, |project, cx| {
196                project
197                    .worktree_for_id(worktree_id, cx)
198                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
199            })??;
200
201            let files = worktree.update(&mut cx, |worktree, _cx| {
202                collect_files_in_path(worktree, &path)
203            })?;
204
205            let open_buffer_tasks = project.update(&mut cx, |project, cx| {
206                files
207                    .into_iter()
208                    .map(|file_path| {
209                        project.open_buffer(
210                            ProjectPath {
211                                worktree_id,
212                                path: file_path.clone(),
213                            },
214                            cx,
215                        )
216                    })
217                    .collect::<Vec<_>>()
218            })?;
219
220            let open_all_buffers_tasks = cx.background_executor().spawn(async move {
221                let mut buffers = Vec::with_capacity(open_buffer_tasks.len());
222
223                for open_buffer_task in open_buffer_tasks {
224                    let buffer = open_buffer_task.await?;
225
226                    buffers.push(buffer);
227                }
228
229                anyhow::Ok(buffers)
230            });
231
232            let buffers = open_all_buffers_tasks.await?;
233
234            this.update(&mut cx, |this, cx| {
235                let mut text = String::new();
236
237                for buffer in buffers {
238                    text.push_str(&codeblock_fence_for_path(Some(&path), None));
239                    text.push_str(&buffer.read(cx).text());
240                    if !text.ends_with('\n') {
241                        text.push('\n');
242                    }
243
244                    text.push_str("```\n");
245                }
246
247                this.delegate
248                    .context_store
249                    .update(cx, |context_store, _cx| {
250                        context_store.insert_context(
251                            ContextKind::Directory,
252                            path.to_string_lossy().to_string(),
253                            text,
254                        );
255                    })?;
256
257                match confirm_behavior {
258                    ConfirmBehavior::KeepOpen => {}
259                    ConfirmBehavior::Close => this.delegate.dismissed(cx),
260                }
261
262                anyhow::Ok(())
263            })??;
264
265            anyhow::Ok(())
266        })
267        .detach_and_log_err(cx)
268    }
269
270    fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
271        self.context_picker
272            .update(cx, |this, cx| {
273                this.reset_mode();
274                cx.emit(DismissEvent);
275            })
276            .ok();
277    }
278
279    fn render_match(
280        &self,
281        ix: usize,
282        selected: bool,
283        _cx: &mut ViewContext<Picker<Self>>,
284    ) -> Option<Self::ListItem> {
285        let path_match = &self.matches[ix];
286        let directory_name = path_match.path.to_string_lossy().to_string();
287
288        Some(
289            ListItem::new(ix)
290                .inset(true)
291                .toggle_state(selected)
292                .child(h_flex().gap_2().child(Label::new(directory_name))),
293        )
294    }
295}
296
297fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
298    let mut files = Vec::new();
299
300    for entry in worktree.child_entries(path) {
301        if entry.is_dir() {
302            files.extend(collect_files_in_path(worktree, &entry.path));
303        } else if entry.is_file() {
304            files.push(entry.path.clone());
305        }
306    }
307
308    files
309}