1use parking_lot::Mutex;
33use serde::{Deserialize, Serialize};
34use std::cmp::Reverse;
35use std::collections::BinaryHeap;
36use std::num::NonZero;
37use std::path::Path;
38use std::sync::Arc;
39use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
40use tokio::sync::RwLock;
41
42use rayon::prelude::*;
43
44pub struct FileIndex {
49 files: Vec<String>,
51 directories: Vec<String>,
53 last_built: std::time::Instant,
55}
56
57fn build_parallel_walker(
59 search_directory: &Path,
60 exclude: &[String],
61 threads: usize,
62 respect_gitignore: bool,
63 follow_links: bool,
64) -> anyhow::Result<ignore::WalkParallel> {
65 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
66 vtcode_commons::walk::apply_defaults(&mut walk_builder);
67
68 walk_builder.threads(threads);
70 walk_builder.follow_links(follow_links);
71 walk_builder.require_git(false); if !respect_gitignore {
74 walk_builder
75 .git_ignore(false)
76 .git_global(false)
77 .git_exclude(false)
78 .ignore(false)
79 .parents(false);
80 }
81
82 if !exclude.is_empty() {
83 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
84 for exclude_pattern in exclude {
85 let pattern = format!("!{exclude_pattern}");
86 override_builder.add(&pattern)?;
87 }
88 walk_builder.overrides(override_builder.build()?);
89 }
90
91 Ok(walk_builder.build_parallel())
92}
93
94impl FileIndex {
95 fn build_from_directory(
98 search_directory: &Path,
99 exclude: &[String],
100 respect_gitignore: bool,
101 threads: usize,
102 ) -> anyhow::Result<Self> {
103 let walker =
104 build_parallel_walker(search_directory, exclude, threads, respect_gitignore, true)?;
105
106 let files_arc = Arc::new(Mutex::new(Vec::new()));
108 let dirs_arc = Arc::new(Mutex::new(Vec::new()));
109
110 walker.run(|| {
111 let files_clone = files_arc.clone();
112 let dirs_clone = dirs_arc.clone();
113 let search_dir = search_directory.to_path_buf();
114
115 Box::new(move |result| {
116 let entry = match result {
117 Ok(e) => e,
118 Err(_) => return ignore::WalkState::Continue,
119 };
120
121 if let Some(rel_path) =
123 entry.path().strip_prefix(&search_dir).ok().and_then(|p| p.to_str())
124 && !rel_path.is_empty()
125 {
126 if entry.path().is_dir() {
127 dirs_clone.lock().push(rel_path.to_string());
128 } else {
129 files_clone.lock().push(rel_path.to_string());
130 }
131 }
132
133 ignore::WalkState::Continue
134 })
135 });
136
137 let files = Arc::try_unwrap(files_arc)
138 .map_err(|arc| {
139 anyhow::anyhow!(
140 "failed to unwrap files arc, {} references remain",
141 Arc::strong_count(&arc)
142 )
143 })?
144 .into_inner();
145 let directories = Arc::try_unwrap(dirs_arc)
146 .map_err(|arc| {
147 anyhow::anyhow!(
148 "failed to unwrap dirs arc, {} references remain",
149 Arc::strong_count(&arc)
150 )
151 })?
152 .into_inner();
153
154 Ok(Self {
155 files,
156 directories,
157 last_built: std::time::Instant::now(),
158 })
159 }
160
161 fn query(
164 &self,
165 pattern_text: &str,
166 limit: usize,
167 match_type_filter: Option<MatchType>,
168 ) -> Vec<(u32, String, MatchType)> {
169 let mut heaps = Vec::new();
174
175 if match_type_filter.is_none_or(|t| t == MatchType::File) {
176 heaps.push(score_paths_top_k(&self.files, limit, pattern_text, MatchType::File));
177 }
178
179 if match_type_filter.is_none_or(|t| t == MatchType::Directory) {
180 heaps.push(score_paths_top_k(
181 &self.directories,
182 limit,
183 pattern_text,
184 MatchType::Directory,
185 ));
186 }
187
188 merge_top_k(heaps, limit)
189 .into_sorted_vec()
190 .into_iter()
191 .map(|Reverse(item)| item)
192 .collect()
193 }
194}
195
196fn score_paths_top_k(
204 paths: &[String],
205 limit: usize,
206 pattern_text: &str,
207 match_type: MatchType,
208) -> BinaryHeap<Reverse<(u32, String, MatchType)>> {
209 const CHUNK: usize = 1024;
210
211 if paths.len() <= CHUNK {
214 let mut list = BestMatchesList::new(limit, pattern_text);
215 for path in paths {
216 list.record_match(path, match_type);
217 }
218 return list.matches;
219 }
220
221 let heaps: Vec<_> = paths
222 .par_chunks(CHUNK)
223 .map_init(
224 || BestMatchesList::new(limit, pattern_text),
225 |list, chunk| {
226 for path in chunk {
227 list.record_match(path, match_type);
228 }
229 std::mem::take(&mut list.matches)
230 },
231 )
232 .collect();
233
234 merge_top_k(heaps, limit)
235}
236
237fn merge_top_k(
243 heaps: Vec<BinaryHeap<Reverse<(u32, String, MatchType)>>>,
244 limit: usize,
245) -> BinaryHeap<Reverse<(u32, String, MatchType)>> {
246 let mut merged = BinaryHeap::with_capacity(limit);
247 for heap in heaps {
248 for Reverse(item) in heap.into_vec() {
249 push_top_match(&mut merged, limit, item.0, item.1, item.2);
250 }
251 }
252 merged
253}
254
255pub struct FileIndexCache {
257 cache: Arc<RwLock<Option<Arc<FileIndex>>>>,
258 search_directory: std::path::PathBuf,
259 exclude: Vec<String>,
260 respect_gitignore: bool,
261 threads: usize,
262}
263
264impl FileIndexCache {
265 pub fn new(
266 search_directory: std::path::PathBuf,
267 exclude: impl IntoIterator<Item = String>,
268 respect_gitignore: bool,
269 threads: usize,
270 ) -> Self {
271 Self {
272 cache: Arc::new(RwLock::new(None)),
273 search_directory,
274 exclude: exclude.into_iter().collect(),
275 respect_gitignore,
276 threads,
277 }
278 }
279
280 pub async fn get_or_build(&self) -> anyhow::Result<Arc<FileIndex>> {
282 {
284 let guard = self.cache.read().await;
285 if let Some(index) = guard.as_ref() {
286 if index.last_built.elapsed() < std::time::Duration::from_secs(300) {
288 return Ok(Arc::clone(index));
289 }
290 }
291 }
292
293 let index = Arc::new(FileIndex::build_from_directory(
295 &self.search_directory,
296 &self.exclude,
297 self.respect_gitignore,
298 self.threads,
299 )?);
300
301 {
303 let mut guard = self.cache.write().await;
304 *guard = Some(Arc::clone(&index));
305 }
306 Ok(index)
307 }
308
309 pub fn refresh_background(&self) -> Option<Arc<FileIndex>> {
312 let search_directory = self.search_directory.clone();
314 let exclude = self.exclude.clone();
315 let respect_gitignore = self.respect_gitignore;
316 let threads = self.threads;
317 let cache = self.cache.clone();
318
319 tokio::spawn(async move {
320 match FileIndex::build_from_directory(
321 &search_directory,
322 &exclude,
323 respect_gitignore,
324 threads,
325 ) {
326 Ok(new_index) => {
327 let mut guard = cache.write().await;
328 *guard = Some(Arc::new(new_index));
329 }
330 Err(e) => {
331 tracing::error!("failed to rebuild file index: {e}");
332 }
333 }
334 });
335
336 let guard = self.cache.blocking_read();
338 guard.as_ref().map(Arc::clone)
339 }
340
341 pub fn update_file(&self, path: &str, is_added: bool) {
344 let mut guard = self.cache.blocking_write();
345 let Some(existing) = guard.take() else { return };
346
347 let mut index = Arc::try_unwrap(existing).unwrap_or_else(|arc| (*arc).clone());
348 if is_added {
349 if Path::new(path).is_dir() {
350 index.directories.push(path.to_string());
351 } else {
352 index.files.push(path.to_string());
353 }
354 } else {
355 index.files.retain(|p| p != path);
356 index.directories.retain(|p| p != path);
357 }
358 index.last_built = std::time::Instant::now();
359 *guard = Some(Arc::new(index));
360 }
361
362 pub async fn index_age(&self) -> Option<std::time::Duration> {
364 let guard = self.cache.read().await;
365 guard.as_ref().map(|idx| idx.last_built.elapsed())
366 }
367}
368
369impl Clone for FileIndex {
371 fn clone(&self) -> Self {
372 Self {
373 files: self.files.clone(),
374 directories: self.directories.clone(),
375 last_built: self.last_built,
376 }
377 }
378}
379
380#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
388#[serde(rename_all = "lowercase")]
389pub enum MatchType {
390 File,
391 Directory,
392}
393
394#[derive(Debug, Clone, Serialize, Deserialize)]
395pub struct FileMatch {
396 pub score: u32,
397 pub path: String,
398 pub match_type: MatchType,
399 #[serde(skip_serializing_if = "Option::is_none")]
400 pub indices: Option<Vec<u32>>,
401}
402
403#[derive(Debug)]
405pub struct FileSearchResults {
406 pub matches: Vec<FileMatch>,
407 pub total_match_count: usize,
408}
409
410pub struct FileSearchConfig {
412 pub pattern_text: String,
413 pub limit: NonZero<usize>,
414 pub search_directory: std::path::PathBuf,
415 pub exclude: Vec<String>,
416 pub threads: NonZero<usize>,
417 pub cancel_flag: Arc<AtomicBool>,
418 pub compute_indices: bool,
419 pub respect_gitignore: bool,
420}
421
422pub use vtcode_commons::paths::file_name_from_path;
423
424struct BestMatchesList {
429 matches: BinaryHeap<Reverse<(u32, String, MatchType)>>,
430 limit: usize,
431 matcher: nucleo_matcher::Matcher,
432 haystack_buf: Vec<char>,
433 pattern: PatternStorage,
435}
436
437enum PatternStorage {
439 Ascii(Vec<u8>),
441 Unicode(Vec<char>),
443}
444
445impl BestMatchesList {
446 fn new(limit: usize, pattern_text: &str) -> Self {
447 let pattern = if pattern_text.is_ascii() {
451 PatternStorage::Ascii(pattern_text.to_ascii_lowercase().into_bytes())
452 } else {
453 PatternStorage::Unicode(pattern_text.to_lowercase().chars().collect())
454 };
455
456 Self {
457 matches: BinaryHeap::new(),
458 limit,
459 matcher: nucleo_matcher::Matcher::new(nucleo_matcher::Config::DEFAULT),
460 haystack_buf: Vec::with_capacity(256),
461 pattern,
462 }
463 }
464
465 fn record_match(&mut self, path: &str, match_type: MatchType) -> bool {
470 let haystack = nucleo_matcher::Utf32Str::new(path, &mut self.haystack_buf);
472 let needle = match &self.pattern {
473 PatternStorage::Ascii(bytes) => nucleo_matcher::Utf32Str::Ascii(bytes),
474 PatternStorage::Unicode(chars) => nucleo_matcher::Utf32Str::Unicode(chars),
475 };
476 let Some(score) = self.matcher.fuzzy_match(haystack, needle) else {
477 return false;
478 };
479
480 push_top_match(&mut self.matches, self.limit, score as u32, path.to_string(), match_type);
481 true
482 }
483}
484
485fn push_top_match(
486 matches: &mut BinaryHeap<Reverse<(u32, String, MatchType)>>,
487 limit: usize,
488 score: u32,
489 path: String,
490 match_type: MatchType,
491) -> bool {
492 let candidate = (score, path, match_type);
493 if matches.len() < limit {
494 matches.push(Reverse(candidate));
495 return true;
496 }
497
498 let Some(minimum) = matches.peek().map(|entry| &entry.0) else {
499 return false;
500 };
501
502 if &candidate <= minimum {
503 return false;
504 }
505
506 matches.pop();
507 matches.push(Reverse(candidate));
508 true
509}
510
511pub async fn run_with_index(
525 config: FileSearchConfig,
526 index_cache: &FileIndexCache,
527) -> anyhow::Result<FileSearchResults> {
528 let limit = config.limit.get();
529 let cancel_flag = &config.cancel_flag;
530 let compute_indices = config.compute_indices;
531
532 let index = index_cache.get_or_build().await?;
534
535 if cancel_flag.load(Ordering::Relaxed) {
537 return Ok(FileSearchResults { matches: Vec::new(), total_match_count: 0 });
538 }
539
540 let matched_paths = tokio::task::spawn_blocking({
543 let pattern_text = config.pattern_text.clone();
544 move || Ok::<_, anyhow::Error>(index.query(&pattern_text, limit, None))
545 })
546 .await??;
547
548 let total_match_count = matched_paths.len();
549
550 let matches = matched_paths
552 .into_iter()
553 .map(|(score, path, match_type)| FileMatch {
554 score,
555 path,
556 match_type,
557 indices: if compute_indices {
558 Some(Vec::new())
559 } else {
560 None
561 },
562 })
563 .collect();
564
565 Ok(FileSearchResults { matches, total_match_count })
566}
567
568pub fn run(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
578 run_with_policy(config, true, false)
579}
580
581pub fn run_bounded_no_follow(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
587 run_bounded_no_follow_with_visit(config, |_| {})
588}
589
590fn run_bounded_no_follow_with_visit(
591 config: FileSearchConfig,
592 mut visit: impl FnMut(&Path),
593) -> anyhow::Result<FileSearchResults> {
594 let limit = config.limit.get();
595 let search_directory = &config.search_directory;
596 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
597 vtcode_commons::walk::apply_defaults(&mut walk_builder);
598 walk_builder
599 .follow_links(false)
600 .require_git(false)
601 .sort_by_file_path(|left, right| left.cmp(right));
602
603 if !config.respect_gitignore {
604 walk_builder
605 .git_ignore(false)
606 .git_global(false)
607 .git_exclude(false)
608 .ignore(false)
609 .parents(false);
610 }
611
612 if !config.exclude.is_empty() {
613 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
614 for exclude_pattern in &config.exclude {
615 override_builder.add(&format!("!{exclude_pattern}"))?;
616 }
617 walk_builder.overrides(override_builder.build()?);
618 }
619
620 let mut matches = BestMatchesList::new(limit, &config.pattern_text);
621 let mut matching_count = 0usize;
622 for result in walk_builder.build() {
623 if config.cancel_flag.load(Ordering::Relaxed) {
624 break;
625 }
626 let entry = match result {
627 Ok(entry) => entry,
628 Err(_) => continue,
629 };
630 visit(entry.path());
631 if !entry.file_type().is_some_and(|file_type| file_type.is_file()) {
632 continue;
633 }
634 let Some(relative_path) = entry
635 .path()
636 .strip_prefix(search_directory)
637 .ok()
638 .and_then(|path| path.to_str())
639 .filter(|path| !path.is_empty())
640 else {
641 continue;
642 };
643 if matches.record_match(relative_path, MatchType::File) {
644 matching_count += 1;
645 if matching_count >= limit {
646 break;
647 }
648 }
649 }
650
651 let matches = matches
652 .matches
653 .into_sorted_vec()
654 .into_iter()
655 .map(|Reverse((score, path, match_type))| FileMatch {
656 score,
657 path,
658 match_type,
659 indices: config.compute_indices.then(Vec::new),
660 })
661 .collect();
662
663 Ok(FileSearchResults {
664 matches,
665 total_match_count: matching_count + usize::from(matching_count >= limit),
668 })
669}
670
671fn run_with_policy(
672 config: FileSearchConfig,
673 follow_links: bool,
674 files_only: bool,
675) -> anyhow::Result<FileSearchResults> {
676 let limit = config.limit.get();
677 let search_directory = &config.search_directory;
678 let exclude = &config.exclude;
679 let threads = config.threads.get();
680 let cancel_flag = &config.cancel_flag;
681 let compute_indices = config.compute_indices;
682 let respect_gitignore = config.respect_gitignore;
683
684 let walker =
685 build_parallel_walker(search_directory, exclude, threads, respect_gitignore, follow_links)?;
686
687 let best_matchers_per_worker: Vec<Arc<Mutex<BestMatchesList>>> = (0..threads)
690 .map(|_| Arc::new(Mutex::new(BestMatchesList::new(limit, &config.pattern_text))))
691 .collect();
692
693 let total_match_count = Arc::new(AtomicUsize::new(0));
694
695 let worker_counter = AtomicUsize::new(0);
698 let worker_count = best_matchers_per_worker.len();
699 walker.run(|| {
700 let worker_id = worker_counter.fetch_add(1, Ordering::Relaxed) % worker_count;
701 let best_list = best_matchers_per_worker[worker_id].clone();
702 let cancel_flag_clone = cancel_flag.clone();
703 let total_match_count_clone = total_match_count.clone();
704
705 Box::new(move |result| {
706 if cancel_flag_clone.load(Ordering::Relaxed) {
708 return ignore::WalkState::Quit;
709 }
710
711 let entry = match result {
712 Ok(e) => e,
713 Err(_) => return ignore::WalkState::Continue,
714 };
715
716 let relative_path =
718 entry.path().strip_prefix(search_directory).ok().and_then(|p| p.to_str());
719
720 let path_to_match = match relative_path {
721 Some(p) if !p.is_empty() => p,
722 _ => return ignore::WalkState::Continue, };
724
725 let match_type = if entry.path().is_dir() {
726 MatchType::Directory
727 } else {
728 MatchType::File
729 };
730
731 if files_only && match_type == MatchType::Directory {
732 return ignore::WalkState::Continue;
733 }
734
735 {
737 let mut list = best_list.lock();
738 if list.record_match(path_to_match, match_type) {
739 total_match_count_clone.fetch_add(1, Ordering::Relaxed);
740 }
741 }
742
743 ignore::WalkState::Continue
744 })
745 });
746
747 let worker_heaps: Vec<BinaryHeap<Reverse<(u32, String, MatchType)>>> = best_matchers_per_worker
749 .into_iter()
750 .map(|arc| std::mem::take(&mut arc.lock().matches))
751 .collect();
752 let merged_matches = merge_top_k(worker_heaps, limit);
753
754 let matches = merged_matches
756 .into_sorted_vec()
757 .into_iter()
758 .map(|Reverse((score, path, match_type))| FileMatch {
759 score,
760 path,
761 match_type,
762 indices: if compute_indices {
763 Some(Vec::new())
764 } else {
765 None
766 },
767 })
768 .collect();
769
770 Ok(FileSearchResults {
771 matches,
772 total_match_count: total_match_count.load(Ordering::Relaxed),
773 })
774}
775
776#[cfg(test)]
777mod tests {
778 use super::{FileSearchConfig, run_bounded_no_follow, run_bounded_no_follow_with_visit};
779 use std::num::NonZero;
780 use std::sync::Arc;
781 use std::sync::atomic::AtomicBool;
782 use tempfile::TempDir;
783
784 fn bounded_paths(workspace: &std::path::Path) -> Vec<String> {
785 run_bounded_no_follow(FileSearchConfig {
786 pattern_text: "widget".to_string(),
787 limit: NonZero::new(2).expect("non-zero limit"),
788 search_directory: workspace.to_path_buf(),
789 exclude: Vec::new(),
790 threads: NonZero::new(4).expect("non-zero threads"),
791 cancel_flag: Arc::new(AtomicBool::new(false)),
792 compute_indices: false,
793 respect_gitignore: true,
794 })
795 .expect("bounded path search")
796 .matches
797 .into_iter()
798 .map(|candidate| candidate.path)
799 .collect()
800 }
801
802 #[test]
803 fn bounded_path_selection_is_stable_across_repeated_walks() {
804 let workspace = TempDir::new().expect("workspace");
805 for directory in ["z", "a", "m", "b", "y"] {
806 let directory = workspace.path().join(directory);
807 std::fs::create_dir(&directory).expect("fixture directory");
808 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n")
809 .expect("fixture source");
810 }
811
812 let expected = bounded_paths(workspace.path());
813 assert_eq!(expected.len(), 2);
814 for _ in 0..20 {
815 assert_eq!(bounded_paths(workspace.path()), expected);
816 }
817 }
818
819 #[test]
820 fn bounded_path_selection_is_the_sorted_prefix_and_stops_early() {
821 let workspace = TempDir::new().expect("workspace");
822 for directory in ["z", "a", "m", "b", "y"] {
823 let directory = workspace.path().join(directory);
824 std::fs::create_dir(&directory).expect("fixture directory");
825 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n")
826 .expect("fixture source");
827 }
828 let mut visited = Vec::new();
829
830 let results = run_bounded_no_follow_with_visit(
831 FileSearchConfig {
832 pattern_text: "widget".to_string(),
833 limit: NonZero::new(2).expect("non-zero limit"),
834 search_directory: workspace.path().to_path_buf(),
835 exclude: Vec::new(),
836 threads: NonZero::new(4).expect("non-zero threads"),
837 cancel_flag: Arc::new(AtomicBool::new(false)),
838 compute_indices: false,
839 respect_gitignore: true,
840 },
841 |path| visited.push(path.to_path_buf()),
842 )
843 .expect("bounded path search");
844 let mut paths =
845 results.matches.into_iter().map(|candidate| candidate.path).collect::<Vec<_>>();
846 paths.sort();
847
848 assert_eq!(paths, vec!["a/widget.rs", "b/widget.rs"]);
849 assert!(
850 visited.len() < 11,
851 "the bounded route must stop before traversing the complete fixture tree"
852 );
853 assert_eq!(results.total_match_count, 3);
854 }
855}