search_command.rs

  1use super::{file_command::FilePlaceholder, SlashCommand, SlashCommandOutput};
  2use anyhow::Result;
  3use assistant_slash_command::SlashCommandOutputSection;
  4use gpui::{AppContext, Task, WeakView};
  5use language::{CodeLabel, HighlightId, LineEnding, LspAdapterDelegate};
  6use semantic_index::SemanticIndex;
  7use std::{
  8    fmt::Write,
  9    path::PathBuf,
 10    sync::{atomic::AtomicBool, Arc},
 11};
 12use ui::{prelude::*, ButtonLike, ElevationIndex, Icon, IconName};
 13use util::ResultExt;
 14use workspace::Workspace;
 15
 16pub(crate) struct SearchSlashCommand;
 17
 18impl SlashCommand for SearchSlashCommand {
 19    fn name(&self) -> String {
 20        "search".into()
 21    }
 22
 23    fn label(&self, cx: &AppContext) -> CodeLabel {
 24        let mut label = CodeLabel::default();
 25        label.push_str("search ", None);
 26        label.push_str(
 27            "--n",
 28            cx.theme().syntax().highlight_id("comment").map(HighlightId),
 29        );
 30        label.filter_range = 0.."search".len();
 31        label
 32    }
 33
 34    fn description(&self) -> String {
 35        "semantic search".into()
 36    }
 37
 38    fn menu_text(&self) -> String {
 39        "Semantic Search".into()
 40    }
 41
 42    fn requires_argument(&self) -> bool {
 43        true
 44    }
 45
 46    fn complete_argument(
 47        &self,
 48        _query: String,
 49        _cancel: Arc<AtomicBool>,
 50        _workspace: WeakView<Workspace>,
 51        _cx: &mut AppContext,
 52    ) -> Task<Result<Vec<String>>> {
 53        Task::ready(Ok(Vec::new()))
 54    }
 55
 56    fn run(
 57        self: Arc<Self>,
 58        argument: Option<&str>,
 59        workspace: WeakView<Workspace>,
 60        _delegate: Arc<dyn LspAdapterDelegate>,
 61        cx: &mut WindowContext,
 62    ) -> Task<Result<SlashCommandOutput>> {
 63        let Some(workspace) = workspace.upgrade() else {
 64            return Task::ready(Err(anyhow::anyhow!("workspace was dropped")));
 65        };
 66        let Some(argument) = argument else {
 67            return Task::ready(Err(anyhow::anyhow!("missing search query")));
 68        };
 69
 70        let mut limit = None;
 71        let mut query = String::new();
 72        for part in argument.split(' ') {
 73            if let Some(parameter) = part.strip_prefix("--") {
 74                if let Ok(count) = parameter.parse::<usize>() {
 75                    limit = Some(count);
 76                    continue;
 77                }
 78            }
 79
 80            query.push_str(part);
 81            query.push(' ');
 82        }
 83        query.pop();
 84
 85        if query.is_empty() {
 86            return Task::ready(Err(anyhow::anyhow!("missing search query")));
 87        }
 88
 89        let project = workspace.read(cx).project().clone();
 90        let fs = project.read(cx).fs().clone();
 91        let project_index =
 92            cx.update_global(|index: &mut SemanticIndex, cx| index.project_index(project, cx));
 93
 94        cx.spawn(|cx| async move {
 95            let results = project_index
 96                .read_with(&cx, |project_index, cx| {
 97                    project_index.search(query.clone(), limit.unwrap_or(5), cx)
 98                })?
 99                .await?;
100
101            let mut loaded_results = Vec::new();
102            for result in results {
103                let (full_path, file_content) =
104                    result.worktree.read_with(&cx, |worktree, _cx| {
105                        let entry_abs_path = worktree.abs_path().join(&result.path);
106                        let mut entry_full_path = PathBuf::from(worktree.root_name());
107                        entry_full_path.push(&result.path);
108                        let file_content = async {
109                            let entry_abs_path = entry_abs_path;
110                            fs.load(&entry_abs_path).await
111                        };
112                        (entry_full_path, file_content)
113                    })?;
114                if let Some(file_content) = file_content.await.log_err() {
115                    loaded_results.push((result, full_path, file_content));
116                }
117            }
118
119            let output = cx
120                .background_executor()
121                .spawn(async move {
122                    let mut text = format!("Search results for {query}:\n");
123                    let mut sections = Vec::new();
124                    for (result, full_path, file_content) in loaded_results {
125                        let range_start = result.range.start.min(file_content.len());
126                        let range_end = result.range.end.min(file_content.len());
127
128                        let start_line =
129                            file_content[0..range_start].matches('\n').count() as u32 + 1;
130                        let end_line = file_content[0..range_end].matches('\n').count() as u32 + 1;
131                        let start_line_byte_offset = file_content[0..range_start]
132                            .rfind('\n')
133                            .map(|pos| pos + 1)
134                            .unwrap_or_default();
135                        let end_line_byte_offset = file_content[range_end..]
136                            .find('\n')
137                            .map(|pos| range_end + pos)
138                            .unwrap_or_else(|| file_content.len());
139
140                        let section_start_ix = text.len();
141                        writeln!(
142                            text,
143                            "```{}:{}-{}",
144                            result.path.display(),
145                            start_line,
146                            end_line,
147                        )
148                        .unwrap();
149                        let mut excerpt =
150                            file_content[start_line_byte_offset..end_line_byte_offset].to_string();
151                        LineEnding::normalize(&mut excerpt);
152                        text.push_str(&excerpt);
153                        writeln!(text, "\n```\n").unwrap();
154                        let section_end_ix = text.len() - 1;
155
156                        sections.push(SlashCommandOutputSection {
157                            range: section_start_ix..section_end_ix,
158                            render_placeholder: Arc::new(move |id, unfold, _| {
159                                FilePlaceholder {
160                                    id,
161                                    path: Some(full_path.clone()),
162                                    line_range: Some(start_line..end_line),
163                                    unfold,
164                                }
165                                .into_any_element()
166                            }),
167                        });
168                    }
169
170                    let query = SharedString::from(query);
171                    sections.push(SlashCommandOutputSection {
172                        range: 0..text.len(),
173                        render_placeholder: Arc::new(move |id, unfold, _cx| {
174                            ButtonLike::new(id)
175                                .style(ButtonStyle::Filled)
176                                .layer(ElevationIndex::ElevatedSurface)
177                                .child(Icon::new(IconName::MagnifyingGlass))
178                                .child(Label::new(query.clone()))
179                                .on_click(move |_, cx| unfold(cx))
180                                .into_any_element()
181                        }),
182                    });
183
184                    SlashCommandOutput { text, sections }
185                })
186                .await;
187
188            Ok(output)
189        })
190    }
191}