From 6a323ce2ddf3d4db98c24c05390b16fd664426b1 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Thu, 24 Feb 2022 15:33:56 +0100 Subject: [PATCH] Implement a basic project-wide search using Aho-Corasick --- crates/project/src/project.rs | 156 +++++++++++++++++++++++++++++++--- 1 file changed, 146 insertions(+), 10 deletions(-) diff --git a/crates/project/src/project.rs b/crates/project/src/project.rs index b77d8add93..02e202110c 100644 --- a/crates/project/src/project.rs +++ b/crates/project/src/project.rs @@ -2055,7 +2055,6 @@ impl Project { if self.is_local() { let (paths_to_search_tx, paths_to_search_rx) = smol::channel::bounded(1024); - // Submit all worktree paths to the queue. let snapshots = self .strong_worktrees(cx) .filter_map(|tree| { @@ -2068,7 +2067,7 @@ impl Project { for (snapshot_abs_path, snapshot) in snapshots { for file in snapshot.files(false, 0) { if paths_to_search_tx - .send((snapshot_abs_path.clone(), file.path.clone())) + .send((snapshot.id(), snapshot_abs_path.clone(), file.path.clone())) .await .is_err() { @@ -2086,12 +2085,12 @@ impl Project { // .ascii_case_insensitive(!case_sensitive) .build(&[&query]), ); - let (matching_paths_tx, matching_paths_rx) = smol::channel::bounded(1024); + let (matching_paths_tx, mut matching_paths_rx) = smol::channel::bounded(1024); + let workers = cx.background().num_cpus(); cx.background() .spawn({ let fs = self.fs.clone(); let background = cx.background().clone(); - let workers = background.num_cpus(); let search = search.clone(); async move { let fs = &fs; @@ -2103,8 +2102,11 @@ impl Project { let mut paths_to_search_rx = paths_to_search_rx.clone(); scope.spawn(async move { let mut path = PathBuf::new(); - while let Some((snapshot_abs_path, file_path)) = - paths_to_search_rx.next().await + while let Some(( + worktree_id, + snapshot_abs_path, + file_path, + )) = paths_to_search_rx.next().await { if matching_paths_tx.is_closed() { break; @@ -2126,7 +2128,7 @@ impl Project { if matches { if matching_paths_tx - .send((snapshot_abs_path, file_path)) + .send((worktree_id, file_path)) .await .is_err() { @@ -2141,10 +2143,70 @@ impl Project { } }) .detach(); - } else { - } - todo!() + let (buffers_tx, buffers_rx) = smol::channel::bounded(1024); + let buffers = self + .buffers_state + .borrow() + .open_buffers + .values() + .filter_map(|b| b.upgrade(cx)) + .collect::>(); + cx.spawn(|this, mut cx| async move { + for buffer in buffers { + let snapshot = buffer.read_with(&cx, |buffer, _| buffer.snapshot()); + buffers_tx.send((buffer, snapshot)).await?; + } + + while let Some(project_path) = matching_paths_rx.next().await { + if let Some(buffer) = this + .update(&mut cx, |this, cx| this.open_buffer(project_path, cx)) + .await + .log_err() + { + let snapshot = buffer.read_with(&cx, |buffer, _| buffer.snapshot()); + buffers_tx.send((buffer, snapshot)).await?; + } + } + + Ok::<_, anyhow::Error>(()) + }) + .detach_and_log_err(cx); + + let background = cx.background().clone(); + cx.background().spawn(async move { + let search = &search; + let mut matched_buffers = Vec::new(); + for _ in 0..workers { + matched_buffers.push(HashMap::default()); + } + background + .scoped(|scope| { + for worker_matched_buffers in matched_buffers.iter_mut() { + let mut buffers_rx = buffers_rx.clone(); + scope.spawn(async move { + while let Some((buffer, snapshot)) = buffers_rx.next().await { + for mat in search.stream_find_iter( + snapshot.as_rope().bytes_in_range(0..snapshot.len()), + ) { + let mat = mat.unwrap(); + let range = snapshot.anchor_before(mat.start()) + ..snapshot.anchor_after(mat.end()); + worker_matched_buffers + .entry(buffer.clone()) + .or_insert(Vec::new()) + .push(range); + } + } + }); + } + }) + .await; + matched_buffers.into_iter().flatten().collect() + }) + } else { + todo!() + } } fn request_lsp( @@ -4814,4 +4876,78 @@ mod tests { "const TWO: usize = one::THREE + one::THREE;" ); } + + #[gpui::test] + async fn test_search(mut cx: gpui::TestAppContext) { + let fs = FakeFs::new(cx.background()); + fs.insert_tree( + "/dir", + json!({ + "one.rs": "const ONE: usize = 1;", + "two.rs": "const TWO: usize = one::ONE + one::ONE;", + "three.rs": "const THREE: usize = one::ONE + two::TWO;", + "four.rs": "const FOUR: usize = one::ONE + three::THREE;", + }), + ) + .await; + let project = Project::test(fs.clone(), &mut cx); + let (tree, _) = project + .update(&mut cx, |project, cx| { + project.find_or_create_local_worktree("/dir", false, cx) + }) + .await + .unwrap(); + let worktree_id = tree.read_with(&cx, |tree, _| tree.id()); + cx.read(|cx| tree.read(cx).as_local().unwrap().scan_complete()) + .await; + + assert_eq!( + search(&project, SearchQuery::Plain("TWO".to_string()), &mut cx).await, + HashMap::from_iter([ + ("two.rs".to_string(), vec![6..9]), + ("three.rs".to_string(), vec![37..40]) + ]) + ); + + let buffer_4 = project + .update(&mut cx, |project, cx| { + project.open_buffer((worktree_id, "four.rs"), cx) + }) + .await + .unwrap(); + buffer_4.update(&mut cx, |buffer, cx| { + buffer.edit([20..28, 31..43], "two::TWO", cx); + }); + + assert_eq!( + search(&project, SearchQuery::Plain("TWO".to_string()), &mut cx).await, + HashMap::from_iter([ + ("two.rs".to_string(), vec![6..9]), + ("three.rs".to_string(), vec![37..40]), + ("four.rs".to_string(), vec![25..28, 36..39]) + ]) + ); + + async fn search( + project: &ModelHandle, + query: SearchQuery, + cx: &mut gpui::TestAppContext, + ) -> HashMap>> { + project + .update(cx, |project, cx| project.search(query, cx)) + .await + .into_iter() + .map(|(buffer, ranges)| { + buffer.read_with(cx, |buffer, _| { + let path = buffer.file().unwrap().path().to_string_lossy().to_string(); + let ranges = ranges + .into_iter() + .map(|range| range.to_offset(buffer)) + .collect::>(); + (path, ranges) + }) + }) + .collect() + } + } }