code_kb_core/
text_scan.rs

1use std::fs::File;
2use std::io::Read;
3use std::path::Path;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Mutex, PoisonError};
6use std::time::{Duration, Instant};
7
8use ignore::{DirEntry, WalkBuilder, WalkState};
9use memchr::{memchr, memchr_iter, memmem};
10
11use crate::workspace::{is_hard_excluded, to_forward_slash};
12
13const MAX_FILE_BYTES: u64 = 2 * 1024 * 1024;
14const BINARY_SNIFF_BYTES: u64 = 8 * 1024;
15const MAX_THREADS: usize = 8;
16
17/// Lines of a project that hold a word, with the coverage of the scan.
18///
19/// A binary file counts as scanned: the scan read it and decided to skip it.
20/// A file over the size cap never reaches the count: the walker drops it.
21#[derive(Debug, Clone)]
22pub struct WordHits {
23    /// `(relative path with '/', 1-based line)`, sorted, one entry per line.
24    pub hits: Vec<(String, usize)>,
25    pub files_scanned: usize,
26    /// Files the walk found, including the ones the budget left unread.
27    pub files_total: usize,
28    /// False when the budget cut the scan.
29    pub complete: bool,
30    pub elapsed: Duration,
31}
32
33/// Finds every line of the project under `root` that holds `word` as a whole word.
34///
35/// Word characters are ASCII letters, digits, `_`, and `$`. The scan is case-sensitive and
36/// follows the same ignore rules as the indexer. Once `budget` is spent, the walk goes on
37/// to count files but reads none.
38pub fn scan_word(root: &Path, word: &str, budget: Duration) -> WordHits {
39    let started = Instant::now();
40    if word.is_empty() {
41        return WordHits {
42            hits: Vec::new(),
43            files_scanned: 0,
44            files_total: 0,
45            complete: true,
46            elapsed: started.elapsed(),
47        };
48    }
49
50    let finder = &memmem::Finder::new(word.as_bytes());
51    let hits = &Mutex::new(Vec::new());
52    let files_total = &AtomicUsize::new(0);
53    let files_scanned = &AtomicUsize::new(0);
54
55    WalkBuilder::new(root)
56        .standard_filters(true)
57        .hidden(false)
58        .max_filesize(Some(MAX_FILE_BYTES))
59        .threads(
60            std::thread::available_parallelism()
61                .map_or(1, usize::from)
62                .min(MAX_THREADS),
63        )
64        .add_custom_ignore_filename(".julieignore")
65        .add_custom_ignore_filename(".code-kb-ignore")
66        .add_custom_ignore_filename(".codekbignore")
67        .filter_entry(|entry| {
68            !is_hard_excluded(&entry.file_name().to_string_lossy()) && !is_nested_repository(entry)
69        })
70        .build_parallel()
71        .run(|| {
72            let mut buf = Vec::new();
73            Box::new(move |entry| {
74                let Ok(entry) = entry else {
75                    return WalkState::Continue;
76                };
77                if !entry.file_type().is_some_and(|t| t.is_file()) {
78                    return WalkState::Continue;
79                }
80                files_total.fetch_add(1, Ordering::Relaxed);
81                if started.elapsed() >= budget {
82                    return WalkState::Continue;
83                }
84                files_scanned.fetch_add(1, Ordering::Relaxed);
85                let Ok(rel) = entry.path().strip_prefix(root) else {
86                    return WalkState::Continue;
87                };
88                if !read_text(entry.path(), &mut buf).unwrap_or(false) {
89                    return WalkState::Continue;
90                }
91                let lines = word_lines(&buf, finder);
92                if !lines.is_empty() {
93                    let rel = to_forward_slash(rel);
94                    hits.lock()
95                        .unwrap_or_else(PoisonError::into_inner)
96                        .extend(lines.into_iter().map(|line| (rel.clone(), line)));
97                }
98                WalkState::Continue
99            })
100        });
101
102    let mut hits = hits
103        .lock()
104        .unwrap_or_else(PoisonError::into_inner)
105        .split_off(0);
106    hits.sort_unstable();
107    let files_total = files_total.load(Ordering::Relaxed);
108    let files_scanned = files_scanned.load(Ordering::Relaxed);
109    WordHits {
110        hits,
111        files_scanned,
112        files_total,
113        complete: files_scanned == files_total,
114        elapsed: started.elapsed(),
115    }
116}
117
118fn is_nested_repository(entry: &DirEntry) -> bool {
119    entry.file_type().is_some_and(|t| t.is_dir()) && entry.path().join(".git").exists()
120}
121
122fn is_word_byte(b: u8) -> bool {
123    b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
124}
125
126/// Fills `buf` with the file, or returns `false` when a NUL byte in the first 8 KiB marks it binary.
127fn read_text(path: &Path, buf: &mut Vec<u8>) -> std::io::Result<bool> {
128    buf.clear();
129    let mut file = File::open(path)?;
130    (&mut file).take(BINARY_SNIFF_BYTES).read_to_end(buf)?;
131    if memchr(0, buf).is_some() {
132        return Ok(false);
133    }
134    file.take(MAX_FILE_BYTES).read_to_end(buf)?;
135    Ok(true)
136}
137
138fn word_lines(buf: &[u8], finder: &memmem::Finder<'_>) -> Vec<usize> {
139    let word_len = finder.needle().len();
140    let mut lines = Vec::new();
141    let mut line = 1;
142    let mut counted_to = 0;
143    for start in finder.find_iter(buf) {
144        let bounded = (start == 0 || !is_word_byte(buf[start - 1]))
145            && buf.get(start + word_len).is_none_or(|&b| !is_word_byte(b));
146        if !bounded {
147            continue;
148        }
149        line += memchr_iter(b'\n', &buf[counted_to..start]).count();
150        counted_to = start;
151        if lines.last() != Some(&line) {
152            lines.push(line);
153        }
154    }
155    lines
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use std::fs;
162    use tempfile::TempDir;
163
164    const AMPLE: Duration = Duration::from_secs(60);
165    const MAX_FILE_BYTES: usize = 2 * 1024 * 1024;
166
167    fn tree(files: &[(&str, &[u8])]) -> TempDir {
168        let dir = crate::safe_tempdir();
169        fs::create_dir(dir.path().join(".git")).unwrap();
170        for (rel, content) in files {
171            let path = dir.path().join(rel);
172            fs::create_dir_all(path.parent().unwrap()).unwrap();
173            fs::write(path, content).unwrap();
174        }
175        dir
176    }
177
178    fn hits(dir: &TempDir, word: &str) -> Vec<(String, usize)> {
179        let found = scan_word(dir.path(), word, AMPLE);
180        assert!(found.complete);
181        found.hits
182    }
183
184    fn at(path: &str, lines: &[usize]) -> Vec<(String, usize)> {
185        lines.iter().map(|&l| (path.to_string(), l)).collect()
186    }
187
188    #[test]
189    fn matches_only_whole_words() {
190        let dir = tree(&[(
191            "a.txt",
192            b"get(x)\nobj.get\ngetter\n_get\n$get\nget_it\na get b\nget1\n1get\n",
193        )]);
194        assert_eq!(hits(&dir, "get"), at("a.txt", &[1, 2, 7]));
195    }
196
197    #[test]
198    fn counts_lines_in_crlf_files() {
199        let dir = tree(&[("a.txt", b"alpha\r\nneedle\r\nbeta\r\nneedle\r\n")]);
200        assert_eq!(hits(&dir, "needle"), at("a.txt", &[2, 4]));
201    }
202
203    #[test]
204    fn finds_a_hit_on_the_last_line_without_a_newline() {
205        let dir = tree(&[("a.txt", b"one\ntwo\nneedle")]);
206        assert_eq!(hits(&dir, "needle"), at("a.txt", &[3]));
207    }
208
209    #[test]
210    fn gives_one_entry_for_two_hits_on_a_line() {
211        let dir = tree(&[("a.txt", b"needle needle\nother\nneedle(needle)\n")]);
212        assert_eq!(hits(&dir, "needle"), at("a.txt", &[1, 3]));
213    }
214
215    #[test]
216    fn follows_gitignore_and_custom_ignore_files() {
217        let dir = tree(&[
218            (".gitignore", b"git_ignored.txt\n"),
219            (".julieignore", b"julie_ignored.txt\n"),
220            (".code-kb-ignore", b"kb_ignored.txt\n"),
221            (".codekbignore", b"kb_alias_ignored.txt\n"),
222            ("git_ignored.txt", b"needle\n"),
223            ("julie_ignored.txt", b"needle\n"),
224            ("kb_ignored.txt", b"needle\n"),
225            ("kb_alias_ignored.txt", b"needle\n"),
226            ("kept.txt", b"needle\n"),
227        ]);
228        assert_eq!(hits(&dir, "needle"), at("kept.txt", &[1]));
229    }
230
231    #[test]
232    fn skips_excluded_folders() {
233        let dir = tree(&[
234            ("node_modules/pkg/index.js", b"needle\n"),
235            ("src/index.js", b"needle\n"),
236        ]);
237        assert_eq!(hits(&dir, "needle"), at("src/index.js", &[1]));
238    }
239
240    #[test]
241    fn skips_nested_repositories_and_worktrees() {
242        let dir = tree(&[
243            ("inner/.git/HEAD", b"ref\n"),
244            ("inner/lib.txt", b"needle\n"),
245            ("linked/.git", b"gitdir: elsewhere\n"),
246            ("linked/lib.txt", b"needle\n"),
247            ("outer.txt", b"needle\n"),
248        ]);
249        assert_eq!(hits(&dir, "needle"), at("outer.txt", &[1]));
250    }
251
252    #[test]
253    fn scans_a_root_whose_own_name_is_excluded() {
254        let parent = crate::safe_tempdir();
255        let root = parent.path().join("build");
256        fs::create_dir_all(&root).unwrap();
257        fs::write(root.join("a.txt"), "needle\n").unwrap();
258        let found = scan_word(&root, "needle", AMPLE);
259        assert_eq!(found.hits, at("a.txt", &[1]));
260    }
261
262    #[test]
263    fn skips_binary_files_by_nul_in_the_first_8_kib() {
264        let mut late_nul = b"needle\n".to_vec();
265        late_nul.extend(vec![b'x'; 9000]);
266        late_nul.push(0);
267        let dir = tree(&[
268            ("early.dat", b"abc\0\nneedle\n"),
269            ("late.dat", late_nul.as_slice()),
270        ]);
271        assert_eq!(hits(&dir, "needle"), at("late.dat", &[1]));
272    }
273
274    #[test]
275    fn skips_files_over_2_mb() {
276        let padding = |extra: usize| {
277            let mut content = b"needle\n".to_vec();
278            content.resize(MAX_FILE_BYTES + extra, b'x');
279            content
280        };
281        let dir = tree(&[
282            ("at_cap.txt", padding(0).as_slice()),
283            ("over_cap.txt", padding(1).as_slice()),
284        ]);
285        assert_eq!(hits(&dir, "needle"), at("at_cap.txt", &[1]));
286    }
287
288    #[test]
289    fn a_zero_budget_reads_no_file_and_reports_the_cut() {
290        let dir = tree(&[
291            ("a.txt", b"needle\n"),
292            ("b.txt", b"needle\n"),
293            ("c/d.txt", b"needle\n"),
294        ]);
295        let found = scan_word(dir.path(), "needle", Duration::ZERO);
296        assert!(!found.complete);
297        assert_eq!(found.files_total, 3);
298        assert_eq!(found.files_scanned, 0);
299        assert!(found.hits.is_empty());
300    }
301
302    #[test]
303    fn a_full_scan_is_complete_with_sorted_slash_paths() {
304        let dir = tree(&[
305            ("b.txt", b"x\nneedle\n"),
306            ("a/b/deep.txt", b"needle\nneedle\nx\nneedle\n"),
307            ("a/top.txt", b"needle\n"),
308        ]);
309        let found = scan_word(dir.path(), "needle", AMPLE);
310        assert!(found.complete);
311        assert_eq!(found.files_scanned, 3);
312        assert_eq!(found.files_total, 3);
313        let mut expected = at("a/b/deep.txt", &[1, 2, 4]);
314        expected.extend(at("a/top.txt", &[1]));
315        expected.extend(at("b.txt", &[2]));
316        assert_eq!(found.hits, expected);
317    }
318
319    #[test]
320    fn an_empty_word_matches_nothing() {
321        let dir = tree(&[("a.txt", b"needle\n")]);
322        let found = scan_word(dir.path(), "", AMPLE);
323        assert!(found.hits.is_empty());
324        assert!(found.complete);
325    }
326}
code_kb_core/text_scan.rs

code_kb_core/
text_scan.rs
