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 let already_included = 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 if already_included {
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 buffers = futures::future::join_all(open_buffer_tasks).await;
235
236 this.update(&mut cx, |this, cx| {
237 let mut text = String::new();
238
239 let mut ok_count = 0;
240
241 for buffer in buffers.into_iter().flatten() {
242 let buffer = buffer.read(cx);
243 let path = buffer.file().map_or(&path, |file| file.path());
244 push_fenced_codeblock(&path, buffer.text(), &mut text);
245 ok_count += 1;
246 }
247
248 if ok_count == 0 {
249 let Some(workspace) = workspace.upgrade() else {
250 return anyhow::Ok(());
251 };
252
253 workspace.update(cx, |workspace, cx| {
254 workspace.show_error(
255 &anyhow::anyhow!(
256 "Could not read any text files from {}",
257 path.display()
258 ),
259 cx,
260 );
261 });
262
263 return anyhow::Ok(());
264 }
265
266 this.delegate
267 .context_store
268 .update(cx, |context_store, _cx| {
269 context_store.insert_directory(&path, text);
270 })?;
271
272 match confirm_behavior {
273 ConfirmBehavior::KeepOpen => {}
274 ConfirmBehavior::Close => this.delegate.dismissed(cx),
275 }
276
277 anyhow::Ok(())
278 })??;
279
280 anyhow::Ok(())
281 })
282 .detach_and_log_err(cx)
283 }
284
285 fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
286 self.context_picker
287 .update(cx, |this, cx| {
288 this.reset_mode();
289 cx.emit(DismissEvent);
290 })
291 .ok();
292 }
293
294 fn render_match(
295 &self,
296 ix: usize,
297 selected: bool,
298 cx: &mut ViewContext<Picker<Self>>,
299 ) -> Option<Self::ListItem> {
300 let path_match = &self.matches[ix];
301 let directory_name = path_match.path.to_string_lossy().to_string();
302
303 let added = self.context_store.upgrade().map_or(false, |context_store| {
304 context_store
305 .read(cx)
306 .included_directory(&path_match.path)
307 .is_some()
308 });
309
310 Some(
311 ListItem::new(ix)
312 .inset(true)
313 .toggle_state(selected)
314 .child(h_flex().gap_2().child(Label::new(directory_name)))
315 .when(added, |el| {
316 el.end_slot(
317 h_flex()
318 .gap_1()
319 .child(
320 Icon::new(IconName::Check)
321 .size(IconSize::Small)
322 .color(Color::Success),
323 )
324 .child(Label::new("Added").size(LabelSize::Small)),
325 )
326 }),
327 )
328 }
329}
330
331fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
332 let mut files = Vec::new();
333
334 for entry in worktree.child_entries(path) {
335 if entry.is_dir() {
336 files.extend(collect_files_in_path(worktree, &entry.path));
337 } else if entry.is_file() {
338 files.push(entry.path.clone());
339 }
340 }
341
342 files
343}