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