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