search_command.rs

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