1use std::cmp::Reverse;
2use std::sync::Arc;
3use std::sync::atomic::AtomicBool;
4
5use anyhow::{Result, anyhow};
6use fuzzy::{StringMatch, StringMatchCandidate};
7use gpui::{
8 App, AppContext, DismissEvent, Entity, FocusHandle, Focusable, Stateful, Task, WeakEntity,
9};
10use ordered_float::OrderedFloat;
11use picker::{Picker, PickerDelegate};
12use project::lsp_store::SymbolLocation;
13use project::{DocumentSymbol, Symbol};
14use ui::{ListItem, prelude::*};
15use util::ResultExt as _;
16use workspace::Workspace;
17
18use crate::context_picker::ContextPicker;
19use agent::context::AgentContextHandle;
20use agent::context_store::ContextStore;
21
22pub struct SymbolContextPicker {
23 picker: Entity<Picker<SymbolContextPickerDelegate>>,
24}
25
26impl SymbolContextPicker {
27 pub fn new(
28 context_picker: WeakEntity<ContextPicker>,
29 workspace: WeakEntity<Workspace>,
30 context_store: WeakEntity<ContextStore>,
31 window: &mut Window,
32 cx: &mut Context<Self>,
33 ) -> Self {
34 let delegate = SymbolContextPickerDelegate::new(context_picker, workspace, context_store);
35 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
36
37 Self { picker }
38 }
39}
40
41impl Focusable for SymbolContextPicker {
42 fn focus_handle(&self, cx: &App) -> FocusHandle {
43 self.picker.focus_handle(cx)
44 }
45}
46
47impl Render for SymbolContextPicker {
48 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
49 self.picker.clone()
50 }
51}
52
53pub struct SymbolContextPickerDelegate {
54 context_picker: WeakEntity<ContextPicker>,
55 workspace: WeakEntity<Workspace>,
56 context_store: WeakEntity<ContextStore>,
57 matches: Vec<SymbolEntry>,
58 selected_index: usize,
59}
60
61impl SymbolContextPickerDelegate {
62 pub fn new(
63 context_picker: WeakEntity<ContextPicker>,
64 workspace: WeakEntity<Workspace>,
65 context_store: WeakEntity<ContextStore>,
66 ) -> Self {
67 Self {
68 context_picker,
69 workspace,
70 context_store,
71 matches: Vec::new(),
72 selected_index: 0,
73 }
74 }
75}
76
77impl PickerDelegate for SymbolContextPickerDelegate {
78 type ListItem = ListItem;
79
80 fn match_count(&self) -> usize {
81 self.matches.len()
82 }
83
84 fn selected_index(&self) -> usize {
85 self.selected_index
86 }
87
88 fn set_selected_index(
89 &mut self,
90 ix: usize,
91 _window: &mut Window,
92 _cx: &mut Context<Picker<Self>>,
93 ) {
94 self.selected_index = ix;
95 }
96
97 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
98 "Search symbols…".into()
99 }
100
101 fn update_matches(
102 &mut self,
103 query: String,
104 window: &mut Window,
105 cx: &mut Context<Picker<Self>>,
106 ) -> Task<()> {
107 let Some(workspace) = self.workspace.upgrade() else {
108 return Task::ready(());
109 };
110
111 let search_task = search_symbols(query, Arc::<AtomicBool>::default(), &workspace, cx);
112 let context_store = self.context_store.clone();
113 cx.spawn_in(window, async move |this, cx| {
114 let symbols = search_task.await;
115
116 let symbol_entries = context_store
117 .read_with(cx, |context_store, cx| {
118 compute_symbol_entries(symbols, context_store, cx)
119 })
120 .log_err()
121 .unwrap_or_default();
122
123 this.update(cx, |this, _cx| {
124 this.delegate.matches = symbol_entries;
125 })
126 .log_err();
127 })
128 }
129
130 fn confirm(&mut self, _secondary: bool, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
131 let Some(mat) = self.matches.get(self.selected_index) else {
132 return;
133 };
134 let Some(workspace) = self.workspace.upgrade() else {
135 return;
136 };
137
138 let add_symbol_task = add_symbol(
139 mat.symbol.clone(),
140 true,
141 workspace,
142 self.context_store.clone(),
143 cx,
144 );
145
146 let selected_index = self.selected_index;
147 cx.spawn(async move |this, cx| {
148 let (_, included) = add_symbol_task.await?;
149 this.update(cx, |this, _| {
150 if let Some(mat) = this.delegate.matches.get_mut(selected_index) {
151 mat.is_included = included;
152 }
153 })
154 })
155 .detach_and_log_err(cx);
156 }
157
158 fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {
159 self.context_picker
160 .update(cx, |_, cx| {
161 cx.emit(DismissEvent);
162 })
163 .ok();
164 }
165
166 fn render_match(
167 &self,
168 ix: usize,
169 selected: bool,
170 _window: &mut Window,
171 _: &mut Context<Picker<Self>>,
172 ) -> Option<Self::ListItem> {
173 let mat = &self.matches.get(ix)?;
174
175 Some(ListItem::new(ix).inset(true).toggle_state(selected).child(
176 render_symbol_context_entry(ElementId::named_usize("symbol-ctx-picker", ix), mat),
177 ))
178 }
179}
180
181pub(crate) struct SymbolEntry {
182 pub symbol: Symbol,
183 pub is_included: bool,
184}
185
186pub(crate) fn add_symbol(
187 symbol: Symbol,
188 remove_if_exists: bool,
189 workspace: Entity<Workspace>,
190 context_store: WeakEntity<ContextStore>,
191 cx: &mut App,
192) -> Task<Result<(Option<AgentContextHandle>, bool)>> {
193 let project = workspace.read(cx).project().clone();
194 let open_buffer_task = project.update(cx, |project, cx| {
195 let SymbolLocation::InProject(symbol_path) = &symbol.path else {
196 return Task::ready(Err(anyhow!("can't add symbol from outside of project")));
197 };
198 project.open_buffer(symbol_path.clone(), cx)
199 });
200 cx.spawn(async move |cx| {
201 let buffer = open_buffer_task.await?;
202 let document_symbols = project
203 .update(cx, |project, cx| project.document_symbols(&buffer, cx))?
204 .await?;
205
206 // Try to find a matching document symbol. Document symbols include
207 // not only the symbol itself (e.g. function name), but they also
208 // include the context that they contain (e.g. function body).
209 let (name, range, enclosing_range) = if let Some(DocumentSymbol {
210 name,
211 range,
212 selection_range,
213 ..
214 }) =
215 find_matching_symbol(&symbol, document_symbols.as_slice())
216 {
217 (name, selection_range, range)
218 } else {
219 // If we do not find a matching document symbol, fall back to
220 // just the symbol itself
221 (symbol.name, symbol.range.clone(), symbol.range)
222 };
223
224 let (range, enclosing_range) = buffer.read_with(cx, |buffer, _| {
225 (
226 buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
227 buffer.anchor_after(enclosing_range.start)
228 ..buffer.anchor_before(enclosing_range.end),
229 )
230 })?;
231
232 context_store.update(cx, move |context_store, cx| {
233 context_store.add_symbol(
234 buffer,
235 name.into(),
236 range,
237 enclosing_range,
238 remove_if_exists,
239 cx,
240 )
241 })
242 })
243}
244
245fn find_matching_symbol(symbol: &Symbol, candidates: &[DocumentSymbol]) -> Option<DocumentSymbol> {
246 let mut candidates = candidates.iter();
247 let mut candidate = candidates.next()?;
248
249 loop {
250 if candidate.range.start > symbol.range.end {
251 return None;
252 }
253 if candidate.range.end < symbol.range.start {
254 candidate = candidates.next()?;
255 continue;
256 }
257 if candidate.selection_range == symbol.range {
258 return Some(candidate.clone());
259 }
260 if candidate.range.start <= symbol.range.start && symbol.range.end <= candidate.range.end {
261 candidates = candidate.children.iter();
262 candidate = candidates.next()?;
263 continue;
264 }
265 return None;
266 }
267}
268
269pub struct SymbolMatch {
270 pub symbol: Symbol,
271}
272
273pub(crate) fn search_symbols(
274 query: String,
275 cancellation_flag: Arc<AtomicBool>,
276 workspace: &Entity<Workspace>,
277 cx: &mut App,
278) -> Task<Vec<SymbolMatch>> {
279 let symbols_task = workspace.update(cx, |workspace, cx| {
280 workspace
281 .project()
282 .update(cx, |project, cx| project.symbols(&query, cx))
283 });
284 let project = workspace.read(cx).project().clone();
285 cx.spawn(async move |cx| {
286 let Some(symbols) = symbols_task.await.log_err() else {
287 return Vec::new();
288 };
289 let Some((visible_match_candidates, external_match_candidates)): Option<(Vec<_>, Vec<_>)> =
290 project
291 .update(cx, |project, cx| {
292 symbols
293 .iter()
294 .enumerate()
295 .map(|(id, symbol)| {
296 StringMatchCandidate::new(id, symbol.label.filter_text())
297 })
298 .partition(|candidate| match &symbols[candidate.id].path {
299 SymbolLocation::InProject(project_path) => project
300 .entry_for_path(project_path, cx)
301 .is_some_and(|e| !e.is_ignored),
302 SymbolLocation::OutsideProject { .. } => false,
303 })
304 })
305 .log_err()
306 else {
307 return Vec::new();
308 };
309
310 const MAX_MATCHES: usize = 100;
311 let mut visible_matches = cx.background_executor().block(fuzzy::match_strings(
312 &visible_match_candidates,
313 &query,
314 false,
315 true,
316 MAX_MATCHES,
317 &cancellation_flag,
318 cx.background_executor().clone(),
319 ));
320 let mut external_matches = cx.background_executor().block(fuzzy::match_strings(
321 &external_match_candidates,
322 &query,
323 false,
324 true,
325 MAX_MATCHES - visible_matches.len().min(MAX_MATCHES),
326 &cancellation_flag,
327 cx.background_executor().clone(),
328 ));
329 let sort_key_for_match = |mat: &StringMatch| {
330 let symbol = &symbols[mat.candidate_id];
331 (Reverse(OrderedFloat(mat.score)), symbol.label.filter_text())
332 };
333
334 visible_matches.sort_unstable_by_key(sort_key_for_match);
335 external_matches.sort_unstable_by_key(sort_key_for_match);
336 let mut matches = visible_matches;
337 matches.append(&mut external_matches);
338
339 matches
340 .into_iter()
341 .map(|mut mat| {
342 let symbol = symbols[mat.candidate_id].clone();
343 let filter_start = symbol.label.filter_range.start;
344 for position in &mut mat.positions {
345 *position += filter_start;
346 }
347 SymbolMatch { symbol }
348 })
349 .collect()
350 })
351}
352
353fn compute_symbol_entries(
354 symbols: Vec<SymbolMatch>,
355 context_store: &ContextStore,
356 cx: &App,
357) -> Vec<SymbolEntry> {
358 symbols
359 .into_iter()
360 .map(|SymbolMatch { symbol, .. }| SymbolEntry {
361 is_included: context_store.includes_symbol(&symbol, cx),
362 symbol,
363 })
364 .collect::<Vec<_>>()
365}
366
367pub fn render_symbol_context_entry(id: ElementId, entry: &SymbolEntry) -> Stateful<Div> {
368 let path = match &entry.symbol.path {
369 SymbolLocation::InProject(project_path) => {
370 project_path.path.file_name().unwrap_or_default().into()
371 }
372 SymbolLocation::OutsideProject {
373 abs_path,
374 signature: _,
375 } => abs_path
376 .file_name()
377 .map(|f| f.to_string_lossy())
378 .unwrap_or_default(),
379 };
380 let symbol_location = format!("{} L{}", path, entry.symbol.range.start.0.row + 1);
381
382 h_flex()
383 .id(id)
384 .gap_1p5()
385 .w_full()
386 .child(
387 Icon::new(IconName::Code)
388 .size(IconSize::Small)
389 .color(Color::Muted),
390 )
391 .child(
392 h_flex()
393 .gap_1()
394 .child(Label::new(&entry.symbol.name))
395 .child(
396 Label::new(symbol_location)
397 .size(LabelSize::Small)
398 .color(Color::Muted),
399 ),
400 )
401 .when(entry.is_included, |el| {
402 el.child(
403 h_flex()
404 .w_full()
405 .justify_end()
406 .gap_0p5()
407 .child(
408 Icon::new(IconName::Check)
409 .size(IconSize::Small)
410 .color(Color::Success),
411 )
412 .child(Label::new("Added").size(LabelSize::Small)),
413 )
414 })
415}