symbol_context_picker.rs

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