_diffctx/edges/similarity/
lexical.rs

1use std::path::Path;
2
3use rayon::prelude::*;
4use rustc_hash::{FxHashMap, FxHashSet};
5
6use crate::config::limits::LEXICAL;
7use crate::config::tokenization::TOKENIZATION;
8use crate::config::weights::{DEFAULT_LANG_WEIGHTS, LANG_WEIGHTS, LangWeights};
9use crate::languages::EXTENSION_TO_LANGUAGE;
10use crate::stopwords::{filter_idents, profile_from_path};
11use crate::types::{Fragment, FragmentId, extract_identifier_list};
12
13use super::super::EdgeDict;
14use super::super::base::EdgeBuilder;
15
16static LANG_ALIAS: &[(&str, &str)] = &[
17    ("bash", "shell"),
18    ("zsh", "shell"),
19    ("fish", "shell"),
20    ("powershell", "shell"),
21];
22
23fn get_lang_weights(path: &Path) -> &LangWeights {
24    let ext = path
25        .extension()
26        .map(|e| format!(".{}", e.to_string_lossy().to_lowercase()))
27        .unwrap_or_default();
28    let lang = EXTENSION_TO_LANGUAGE.get(ext.as_str()).copied();
29    if let Some(lang) = lang {
30        let aliased = LANG_ALIAS
31            .iter()
32            .find(|(k, _)| *k == lang)
33            .map(|(_, v)| *v)
34            .unwrap_or(lang);
35        LANG_WEIGHTS.get(aliased).unwrap_or(&DEFAULT_LANG_WEIGHTS)
36    } else {
37        &DEFAULT_LANG_WEIGHTS
38    }
39}
40
41fn clamp_lexical_weight(raw_sim: f64, src_path: Option<&Path>, dst_path: Option<&Path>) -> f64 {
42    let (lex_max, lex_min) = match (src_path, dst_path) {
43        (Some(sp), Some(dp)) => {
44            let sw = get_lang_weights(sp);
45            let dw = get_lang_weights(dp);
46            (
47                sw.lexical_max.max(dw.lexical_max),
48                sw.lexical_min.max(dw.lexical_min),
49            )
50        }
51        _ => (LEXICAL.weight_max, LEXICAL.weight_min),
52    };
53
54    if raw_sim < LEXICAL.min_similarity {
55        return 0.0;
56    }
57    let denom = 1.0 - LEXICAL.min_similarity;
58    if denom <= 0.0 {
59        return lex_max;
60    }
61    let normalized = (raw_sim - LEXICAL.min_similarity) / denom;
62    lex_min + normalized * (lex_max - lex_min)
63}
64
65pub struct LexicalEdgeBuilder;
66
67/// Maps each unique term to a compact u32 id. Stores each term string exactly once.
68struct TermInterner {
69    by_str: FxHashMap<String, u32>,
70}
71
72impl TermInterner {
73    fn new() -> Self {
74        Self {
75            by_str: FxHashMap::default(),
76        }
77    }
78
79    fn intern(&mut self, term: String) -> u32 {
80        let next_id = self.by_str.len() as u32;
81        *self.by_str.entry(term).or_insert(next_id)
82    }
83
84    fn len(&self) -> usize {
85        self.by_str.len()
86    }
87}
88
89impl LexicalEdgeBuilder {
90    /// Tokenize and filter identifiers for one fragment. Returns the raw filtered identifier list.
91    fn tokens(frag: &Fragment) -> Vec<String> {
92        let profile = profile_from_path(frag.path());
93        let idents =
94            extract_identifier_list(&frag.content, TOKENIZATION.query_min_identifier_length);
95        filter_idents(&idents, 3, profile)
96    }
97}
98
99impl EdgeBuilder for LexicalEdgeBuilder {
100    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
101        if fragments.is_empty() {
102            return FxHashMap::default();
103        }
104
105        let n_docs = fragments.len();
106        let max_df = (n_docs as f64 * LEXICAL.max_df_ratio).max(1.0) as usize;
107
108        // Pass 1: tokenize each fragment in parallel; flatten to per-fragment Vec<String>.
109        let per_frag_tokens: Vec<Vec<String>> =
110            fragments.par_iter().map(|f| Self::tokens(f)).collect();
111
112        // Pass 2: build the term interner serially, computing document frequency in one go.
113        let mut interner = TermInterner::new();
114        let mut doc_freq: Vec<u32> = Vec::new();
115        let per_frag_term_ids: Vec<Vec<u32>> = per_frag_tokens
116            .into_iter()
117            .map(|tokens| {
118                let mut seen_in_doc: FxHashSet<u32> = FxHashSet::default();
119                let mut ids: Vec<u32> = Vec::with_capacity(tokens.len());
120                for tok in tokens {
121                    let id = interner.intern(tok);
122                    if doc_freq.len() <= id as usize {
123                        doc_freq.resize(id as usize + 1, 0);
124                    }
125                    if seen_in_doc.insert(id) {
126                        doc_freq[id as usize] += 1;
127                    }
128                    ids.push(id);
129                }
130                ids
131            })
132            .collect();
133
134        let n_terms = interner.len();
135        // Interner string-table is no longer needed once doc_freq has been built.
136        drop(interner);
137
138        let n_docs_f = n_docs as f64;
139        let mut idf: Vec<f32> = Vec::with_capacity(n_terms);
140        for &df in &doc_freq {
141            let v = ((n_docs_f + 1.0) / (df as f64 + 1.0)).ln() + 1.0;
142            idf.push(v as f32);
143        }
144
145        // Pass 3: build TF-IDF vectors as sparse Vec<(TermId, f32)>, normalized.
146        let tf_idf: Vec<Vec<(u32, f32)>> = per_frag_term_ids
147            .par_iter()
148            .map(|term_ids| {
149                let mut tf: FxHashMap<u32, u32> = FxHashMap::default();
150                for &id in term_ids {
151                    *tf.entry(id).or_insert(0) += 1;
152                }
153                let mut vec: Vec<(u32, f32)> = Vec::with_capacity(tf.len());
154                for (&term_id, &count) in &tf {
155                    let df = doc_freq[term_id as usize] as usize;
156                    if df == 0 || df > max_df {
157                        continue;
158                    }
159                    let term_idf = idf[term_id as usize];
160                    if (term_idf as f64) < LEXICAL.min_idf {
161                        continue;
162                    }
163                    vec.push((term_id, count as f32 * term_idf));
164                }
165                let norm: f32 = vec.iter().map(|(_, w)| w * w).sum::<f32>().sqrt();
166                if norm > 0.0 {
167                    for (_, w) in &mut vec {
168                        *w /= norm;
169                    }
170                }
171                vec.sort_unstable_by_key(|&(id, _)| id);
172                vec
173            })
174            .collect();
175
176        drop(per_frag_term_ids);
177        drop(doc_freq);
178        drop(idf);
179
180        // Pass 4: invert into postings — for each term, list of (frag_idx, weight).
181        // Consume tf_idf as we go so it never coexists with the inverted index.
182        let mut postings: Vec<Vec<(u32, f32)>> = vec![Vec::new(); n_terms];
183        for (frag_idx, vec) in tf_idf.into_iter().enumerate() {
184            for (term_id, weight) in vec {
185                postings[term_id as usize].push((frag_idx as u32, weight));
186            }
187        }
188
189        // Pass 5: O(F²) inner loop over each posting, capped by max_postings.
190        // Drop each posting list as soon as we are done with it.
191        let mut dot_products: FxHashMap<(u32, u32), f32> = FxHashMap::default();
192        for posting_list in postings.iter_mut() {
193            if posting_list.len() > LEXICAL.max_postings || posting_list.len() < 2 {
194                posting_list.clear();
195                posting_list.shrink_to_fit();
196                continue;
197            }
198            for i in 0..posting_list.len() {
199                let (frag_i, weight_i) = posting_list[i];
200                for j in (i + 1)..posting_list.len() {
201                    let (frag_j, weight_j) = posting_list[j];
202                    let pair = if frag_i < frag_j {
203                        (frag_i, frag_j)
204                    } else {
205                        (frag_j, frag_i)
206                    };
207                    *dot_products.entry(pair).or_insert(0.0) += weight_i * weight_j;
208                }
209            }
210            posting_list.clear();
211            posting_list.shrink_to_fit();
212        }
213        drop(postings);
214
215        // Pass 6: turn pairwise similarities into per-node top-k candidate edges.
216        let frag_paths: Vec<&str> = fragments.iter().map(|f| f.path()).collect();
217        let mut neighbors_by_node: FxHashMap<u32, Vec<(f32, u32)>> = FxHashMap::default();
218
219        let min_sim = LEXICAL.min_similarity as f32;
220        let backward_factor = LEXICAL.backward_factor as f32;
221        for ((src_idx, dst_idx), sim) in &dot_products {
222            if *sim < min_sim {
223                continue;
224            }
225            let src_path = Path::new(frag_paths[*src_idx as usize]);
226            let dst_path = Path::new(frag_paths[*dst_idx as usize]);
227            let fwd = clamp_lexical_weight(*sim as f64, Some(src_path), Some(dst_path)) as f32;
228            let bwd = clamp_lexical_weight(*sim as f64, Some(dst_path), Some(src_path)) as f32
229                * backward_factor;
230            neighbors_by_node
231                .entry(*src_idx)
232                .or_default()
233                .push((fwd, *dst_idx));
234            neighbors_by_node
235                .entry(*dst_idx)
236                .or_default()
237                .push((bwd, *src_idx));
238        }
239
240        let frag_ids: Vec<&FragmentId> = fragments.iter().map(|f| &f.id).collect();
241        let mut edges: EdgeDict = FxHashMap::default();
242        for (node_idx, mut candidates) in neighbors_by_node {
243            candidates.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
244            candidates.truncate(LEXICAL.top_k_neighbors);
245            for (weight, neighbor_idx) in candidates {
246                let key = (
247                    frag_ids[node_idx as usize].clone(),
248                    frag_ids[neighbor_idx as usize].clone(),
249                );
250                let existing = edges.get(&key).copied().unwrap_or(0.0);
251                let weight_f64 = weight as f64;
252                if weight_f64 > existing {
253                    edges.insert(key, weight_f64);
254                }
255            }
256        }
257
258        edges
259    }
260
261    fn is_expensive(&self) -> bool {
262        true
263    }
264}
_diffctx/edges/similarity/lexical.rs

_diffctx/edges/similarity/
lexical.rs
