_diffctx/
token_corpus.rs

1use std::path::{Path, PathBuf};
2use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
3use std::time::Instant;
4
5use rayon::prelude::*;
6use rustc_hash::{FxHashMap, FxHashSet};
7use serde::{Deserialize, Serialize};
8
9use crate::config::bm25::BM25;
10use crate::config::tokenization::TOKENIZATION;
11use crate::discovery::DiscoveryContext;
12use crate::git;
13
14// Bump whenever identifier extraction changes (regex, lowercasing, length
15// filtering): the epoch is part of the on-disk cache key, so stale entries
16// are invalidated instead of silently poisoning results.
17pub const TOKENIZER_EPOCH: u32 = 1;
18
19pub struct DocTokens {
20    pub term_counts: FxHashMap<String, u32>,
21    pub total_len: u32,
22}
23
24pub struct TokenCorpus {
25    pub docs: Vec<(PathBuf, DocTokens)>,
26}
27
28// The rare-identifier and BM25 strategies can share one tokenized corpus
29// only because their tokenizers are identical: same identifier regex, same
30// lowercasing, equal minimum lengths. The assert pins that assumption -
31// filtering by post-lowercase length is NOT equivalent to the pre-lowercase
32// length filter for non-ASCII identifiers, so if these configs ever diverge
33// the corpus must be split back into per-strategy passes.
34fn shared_min_token_length() -> usize {
35    debug_assert_eq!(
36        TOKENIZATION.query_min_identifier_length,
37        BM25.min_query_token_length
38    );
39    TOKENIZATION
40        .query_min_identifier_length
41        .min(BM25.min_query_token_length)
42}
43
44impl TokenCorpus {
45    pub fn build(ctx: &DiscoveryContext) -> Self {
46        let t0 = Instant::now();
47        let min_len = shared_min_token_length();
48        let changed_set: FxHashSet<&Path> = ctx.changed_files.iter().map(|p| p.as_path()).collect();
49        let store = TokenCacheStore::open(min_len);
50        let oids = if store.is_some() {
51            resolve_cacheable_oids(&ctx.root_dir)
52        } else {
53            FxHashMap::default()
54        };
55
56        let hits = AtomicUsize::new(0);
57        let tokenized = AtomicUsize::new(0);
58        let docs: Vec<(PathBuf, DocTokens)> = ctx
59            .all_candidates
60            .par_iter()
61            .filter(|f| !changed_set.contains(f.as_path()))
62            .filter_map(|f| {
63                let oid = oids.get(f.as_path());
64                if let (Some(store), Some(oid)) = (store.as_ref(), oid) {
65                    if let Some(doc) = store.load(oid) {
66                        hits.fetch_add(1, Ordering::Relaxed);
67                        return Some((f.clone(), doc));
68                    }
69                }
70                let content = ctx.read_file(f)?;
71                let (term_counts, total_len) =
72                    crate::types::extract_identifier_counts(&content, min_len);
73                let doc = DocTokens {
74                    term_counts,
75                    total_len,
76                };
77                if let (Some(store), Some(oid)) = (store.as_ref(), oid) {
78                    store.save(oid, &doc);
79                }
80                tokenized.fetch_add(1, Ordering::Relaxed);
81                Some((f.clone(), doc))
82            })
83            .collect();
84
85        tracing::debug!(
86            "token corpus: {} docs ({} cache hits, {} tokenized) in {:.3}s",
87            docs.len(),
88            hits.load(Ordering::Relaxed),
89            tokenized.load(Ordering::Relaxed),
90            t0.elapsed().as_secs_f64(),
91        );
92        Self { docs }
93    }
94}
95
96// Blob OIDs are only usable as cache keys for regular tracked files whose
97// working-tree content matches the index: symlinks/gitlinks, conflicted
98// stages, and files with unstaged modifications all bypass the cache and
99// fall through to a direct read+tokenize, which keeps cold and warm runs
100// bit-equivalent. Any git failure disables keying entirely for the run.
101fn resolve_cacheable_oids(root_dir: &Path) -> FxHashMap<PathBuf, String> {
102    let Ok(entries) = git::run_git_z(root_dir, &["ls-files", "-s", "-z"]) else {
103        return FxHashMap::default();
104    };
105    let Ok(dirty_parts) = git::run_git_z(root_dir, &["diff-files", "--name-only", "-z"]) else {
106        return FxHashMap::default();
107    };
108    let dirty: FxHashSet<PathBuf> = dirty_parts.into_iter().map(|p| root_dir.join(p)).collect();
109
110    let mut oids: FxHashMap<PathBuf, String> = FxHashMap::default();
111    for entry in entries {
112        let Some((meta, rel)) = entry.split_once('\t') else {
113            continue;
114        };
115        let mut fields = meta.split_ascii_whitespace();
116        let (Some(mode), Some(oid), Some(stage)) = (fields.next(), fields.next(), fields.next())
117        else {
118            continue;
119        };
120        if stage != "0" || !(mode == "100644" || mode == "100755") {
121            continue;
122        }
123        let path = root_dir.join(rel);
124        if dirty.contains(&path) {
125            continue;
126        }
127        oids.insert(path, oid.to_string());
128    }
129    oids
130}
131
132#[derive(Serialize, Deserialize)]
133struct StoredDoc {
134    len: u32,
135    terms: Vec<(String, u32)>,
136}
137
138struct TokenCacheStore {
139    dir: PathBuf,
140}
141
142static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
143
144impl TokenCacheStore {
145    fn open(min_token_length: usize) -> Option<Self> {
146        let root = std::env::var_os("DIFFCTX_TOKEN_CACHE_DIR")
147            .filter(|v| !v.is_empty())
148            .map(PathBuf::from)
149            .or_else(default_cache_root)?;
150        let dir = root.join(format!("v{TOKENIZER_EPOCH}-l{min_token_length}"));
151        std::fs::create_dir_all(&dir).ok()?;
152        Some(Self { dir })
153    }
154
155    fn entry_path(&self, oid: &str) -> Option<PathBuf> {
156        if oid.len() < 3 || !oid.bytes().all(|b| b.is_ascii_hexdigit()) {
157            return None;
158        }
159        Some(self.dir.join(&oid[..2]).join(&oid[2..]))
160    }
161
162    fn load(&self, oid: &str) -> Option<DocTokens> {
163        let bytes = std::fs::read(self.entry_path(oid)?).ok()?;
164        let stored: StoredDoc = serde_json::from_slice(&bytes).ok()?;
165        Some(DocTokens {
166            term_counts: stored.terms.into_iter().collect(),
167            total_len: stored.len,
168        })
169    }
170
171    fn save(&self, oid: &str, doc: &DocTokens) {
172        let Some(path) = self.entry_path(oid) else {
173            return;
174        };
175        let mut terms: Vec<(String, u32)> = doc
176            .term_counts
177            .iter()
178            .map(|(t, c)| (t.clone(), *c))
179            .collect();
180        terms.sort();
181        let stored = StoredDoc {
182            len: doc.total_len,
183            terms,
184        };
185        let Ok(bytes) = serde_json::to_vec(&stored) else {
186            return;
187        };
188        let Some(parent) = path.parent() else {
189            return;
190        };
191        if std::fs::create_dir_all(parent).is_err() {
192            return;
193        }
194        let tmp = parent.join(format!(
195            ".{}.{}.{}.tmp",
196            &oid[2..],
197            std::process::id(),
198            TMP_COUNTER.fetch_add(1, Ordering::Relaxed),
199        ));
200        if std::fs::write(&tmp, &bytes).is_err() {
201            let _ = std::fs::remove_file(&tmp);
202            return;
203        }
204        if std::fs::rename(&tmp, &path).is_err() {
205            let _ = std::fs::remove_file(&tmp);
206        }
207    }
208}
209
210fn default_cache_root() -> Option<PathBuf> {
211    #[cfg(target_os = "macos")]
212    {
213        std::env::var_os("HOME")
214            .map(|h| PathBuf::from(h).join("Library/Caches/diffctx/token-cache"))
215    }
216    #[cfg(target_os = "windows")]
217    {
218        std::env::var_os("LOCALAPPDATA").map(|d| PathBuf::from(d).join("diffctx/token-cache"))
219    }
220    #[cfg(not(any(target_os = "macos", target_os = "windows")))]
221    {
222        std::env::var_os("XDG_CACHE_HOME")
223            .filter(|v| !v.is_empty())
224            .map(PathBuf::from)
225            .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache")))
226            .map(|c| c.join("diffctx/token-cache"))
227    }
228}
_diffctx/token_corpus.rs

_diffctx/
token_corpus.rs
