Skip to main content

ctx/
fingerprint.rs

1//! Structural fingerprinting for near-duplicate function detection.
2//!
3//! This module implements a MinHash-based pipeline over normalized token
4//! shingles:
5//!
6//! 1. **Tokenize** a whole source file with tree-sitter and normalize the
7//!    leaf tokens: identifiers become `ID`, string/number literals become
8//!    `LIT`, comments are dropped, and keywords/operators/structure tokens
9//!    are kept verbatim.
10//! 2. **Shingle** each function/method symbol's token stream into k=5 token
11//!    windows, hashed with FNV-1a (fixed seeds, little-endian byte order, so
12//!    hashes are deterministic across runs and platforms).
13//! 3. **MinHash** the shingle set with 128 permutations
14//!    (`h_i(x) = splitmix64(x ^ SEED[i])`) into a 1024-byte signature stored
15//!    in the `symbol_fingerprints` table during `ctx index`.
16//! 4. **LSH banding** (16 bands x 8 rows) over the signatures produces
17//!    candidate pairs at query time, which are then verified with the exact
18//!    Jaccard similarity over re-derived shingle sets.
19//!
20//! Notes and limitations:
21//! - Solidity has no tree-sitter grammar in this build; it is tokenized with
22//!   the solang-parser lexer instead (see `tokenize`), so Solidity functions
23//!   are fingerprinted just like the tree-sitter languages.
24//! - Nested functions share tokens with their enclosing function: both are
25//!   fingerprinted over their own line ranges, and the parent's range
26//!   includes the child's tokens.
27//! - Idiomatic boilerplate (builders, trait impls, small CRUD handlers) can
28//!   legitimately look structurally similar; tune `--min-tokens` to filter
29//!   short functions.
30
31use std::cell::RefCell;
32use std::collections::hash_map::Entry;
33use std::collections::{HashMap, HashSet};
34use std::path::Path;
35
36use tree_sitter::Parser;
37
38use crate::db::{Database, Fingerprint, Symbol, SymbolKind};
39use crate::error::Result;
40use crate::parser::Language;
41
42/// Number of MinHash permutations (signature length in u64 words).
43pub const NUM_PERMS: usize = 128;
44
45/// Shingle size in tokens.
46pub const SHINGLE_K: usize = 5;
47
48/// Number of LSH bands.
49pub const LSH_BANDS: usize = 16;
50
51/// Rows (signature words) per LSH band.
52pub const LSH_ROWS: usize = 8;
53
54/// Minimum usable similarity threshold. Below this the 16x8 LSH banding
55/// misses too many candidate pairs to be trustworthy, so callers clamp to it.
56pub const MIN_THRESHOLD: f64 = 0.5;
57
58const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
59const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
60
61/// Fixed seed for the SplitMix64 stream that generates the permutation seeds.
62const SEED_STATE: u64 = 0x0c7f_1d5e_a5ed_c0de;
63
64/// SplitMix64 mixing function (Steele, Lea & Flood). Deterministic and
65/// platform-independent; used both to derive the permutation seeds and as
66/// the per-permutation hash `h_i(x) = splitmix64(x ^ SEED[i])`.
67const fn splitmix64(mut z: u64) -> u64 {
68    z = z.wrapping_add(0x9e37_79b9_7f4a_7c15);
69    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
70    z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
71    z ^ (z >> 31)
72}
73
74/// Per-permutation seeds, generated by a SplitMix64 stream from a fixed
75/// constant at compile time.
76const SEEDS: [u64; NUM_PERMS] = {
77    let mut seeds = [0u64; NUM_PERMS];
78    let mut state = SEED_STATE;
79    let mut i = 0;
80    while i < NUM_PERMS {
81        state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
82        seeds[i] = splitmix64(state);
83        i += 1;
84    }
85    seeds
86};
87
88/// FNV-1a 64-bit hash.
89fn fnv1a(bytes: &[u8]) -> u64 {
90    let mut h = FNV_OFFSET;
91    for &b in bytes {
92        h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
93    }
94    h
95}
96
97/// A normalized token: the normalized text plus the 1-indexed source line.
98#[derive(Debug, Clone, PartialEq, Eq)]
99pub struct Tok {
100    pub text: String,
101    pub line: u32,
102}
103
104/// Map a [`Language`] to its tree-sitter grammar, or `None` for languages
105/// without one. Solidity has no grammar here and returns `None`; it is
106/// tokenized with the solang-parser lexer in [`tokenize`] instead.
107fn ts_language(lang: Language) -> Option<tree_sitter::Language> {
108    match lang {
109        Language::Rust => Some(tree_sitter_rust::language()),
110        Language::TypeScript => Some(tree_sitter_typescript::language_typescript()),
111        Language::Tsx => Some(tree_sitter_typescript::language_tsx()),
112        Language::JavaScript | Language::Jsx => Some(tree_sitter_javascript::language()),
113        Language::Python => Some(tree_sitter_python::language()),
114        Language::Go => Some(tree_sitter_go::language()),
115        Language::Solidity | Language::Yaml | Language::Unknown => None,
116    }
117}
118
119thread_local! {
120    /// Per-thread parser cache (indexing is rayon-parallel; tree-sitter
121    /// parsers are not Sync). Keyed by language name.
122    static PARSERS: RefCell<HashMap<&'static str, Parser>> = RefCell::new(HashMap::new());
123}
124
125/// Normalize a leaf node into its token text.
126fn normalize_leaf(kind: &str, text: &str) -> String {
127    if kind.contains("identifier") {
128        return "ID".to_string();
129    }
130    if kind.contains("string")
131        || kind.contains("char")
132        || kind.contains("template")
133        || kind.contains("rune")
134    {
135        return "LIT".to_string();
136    }
137    // Numeric literal node kinds across the supported grammars
138    // (rust, typescript/javascript, python, go).
139    if matches!(
140        kind,
141        "integer_literal"
142            | "float_literal"
143            | "number"
144            | "int_literal"
145            | "imaginary_literal"
146            | "integer"
147            | "float"
148    ) {
149        return "LIT".to_string();
150    }
151    text.to_string()
152}
153
154/// Tokenize a whole source file into a normalized token stream.
155///
156/// Solidity is lexed with the solang-parser lexer; every other supported
157/// language is parsed with tree-sitter and reduced to its leaf nodes in
158/// document order. Comments are dropped in both paths. Returns `None` for
159/// languages without either backend (YAML, unknown) or when parsing fails.
160pub fn tokenize(lang: Language, source: &str) -> Option<Vec<Tok>> {
161    if lang == Language::Solidity {
162        return Some(tokenize_solidity(source));
163    }
164
165    let ts_lang = ts_language(lang)?;
166
167    PARSERS.with(|cell| {
168        let mut parsers = cell.borrow_mut();
169        let parser = parsers.entry(lang.as_str()).or_insert_with(|| {
170            let mut p = Parser::new();
171            p.set_language(ts_lang)
172                .expect("grammar/version mismatch for fingerprint parser");
173            p
174        });
175
176        let tree = parser.parse(source, None)?;
177        let bytes = source.as_bytes();
178        let mut tokens = Vec::new();
179
180        // Iterative DFS (explicit stack) to avoid recursion depth limits on
181        // deeply nested code. Children are pushed in reverse so tokens come
182        // out in document order.
183        let mut stack = vec![tree.root_node()];
184        while let Some(node) = stack.pop() {
185            let kind = node.kind();
186            if kind.contains("comment") {
187                continue;
188            }
189            if node.child_count() == 0 {
190                let range = node.byte_range();
191                if range.is_empty() {
192                    continue; // zero-width "missing" nodes
193                }
194                let text = std::str::from_utf8(&bytes[range]).unwrap_or("");
195                tokens.push(Tok {
196                    text: normalize_leaf(kind, text),
197                    line: node.start_position().row as u32 + 1,
198                });
199                continue;
200            }
201            for i in (0..node.child_count()).rev() {
202                if let Some(child) = node.child(i) {
203                    stack.push(child);
204                }
205            }
206        }
207
208        Some(tokens)
209    })
210}
211
212/// Tokenize Solidity source with the solang-parser lexer.
213///
214/// The lexer yields `(start, token, end)` spans with comments already
215/// excluded, so normalization mirrors [`normalize_leaf`]: identifiers become
216/// `ID`; string/hex/address/number literals become `LIT`; keywords and
217/// punctuation are kept as their verbatim lexeme (via the token's `Display`).
218/// Each token's 1-indexed line is derived from its start byte offset, matching
219/// the line numbers the solang parser records for symbols.
220fn tokenize_solidity(source: &str) -> Vec<Tok> {
221    use solang_parser::lexer::{Lexer, Token};
222
223    let line_starts = line_start_offsets(source);
224
225    // The lexer collects comments and lexical errors out-of-band; we only
226    // consume the token stream, so both sinks are discarded.
227    let mut comments = Vec::new();
228    let mut errors = Vec::new();
229    let lexer = Lexer::new(source, 0, &mut comments, &mut errors);
230
231    lexer
232        .map(|(start, token, _end)| {
233            let text = match token {
234                Token::Identifier(_) => "ID".to_string(),
235                Token::StringLiteral(..)
236                | Token::AddressLiteral(_)
237                | Token::HexLiteral(_)
238                | Token::Number(..)
239                | Token::RationalNumber(..)
240                | Token::HexNumber(_) => "LIT".to_string(),
241                other => other.to_string(),
242            };
243            Tok {
244                text,
245                line: offset_to_line(&line_starts, start),
246            }
247        })
248        .collect()
249}
250
251/// Byte offset of the first character of each source line (line 1 at index 0).
252fn line_start_offsets(source: &str) -> Vec<usize> {
253    let mut starts = vec![0usize];
254    for (i, b) in source.bytes().enumerate() {
255        if b == b'\n' {
256            starts.push(i + 1);
257        }
258    }
259    starts
260}
261
262/// Map a byte offset to its 1-indexed source line via the precomputed
263/// line-start table (the count of line starts at or before the offset).
264fn offset_to_line(line_starts: &[usize], offset: usize) -> u32 {
265    line_starts.partition_point(|&s| s <= offset) as u32
266}
267
268/// Build the set of k-shingle hashes for a token stream.
269///
270/// Each token is hashed with FNV-1a; each shingle hash is FNV-1a over the
271/// k token hashes' little-endian bytes. Streams shorter than
272/// [`SHINGLE_K`] tokens produce an empty set.
273pub fn shingle_set(tokens: &[Tok]) -> HashSet<u64> {
274    let mut shingles = HashSet::new();
275    if tokens.len() < SHINGLE_K {
276        return shingles;
277    }
278    let hashes: Vec<u64> = tokens.iter().map(|t| fnv1a(t.text.as_bytes())).collect();
279    for window in hashes.windows(SHINGLE_K) {
280        let mut h = FNV_OFFSET;
281        for th in window {
282            for b in th.to_le_bytes() {
283                h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
284            }
285        }
286        shingles.insert(h);
287    }
288    shingles
289}
290
291/// Compute the 128-permutation MinHash signature of a shingle set.
292///
293/// Returns `None` for an empty set (no signature can be defined).
294pub fn minhash(shingles: &HashSet<u64>) -> Option<[u64; NUM_PERMS]> {
295    if shingles.is_empty() {
296        return None;
297    }
298    let mut sig = [u64::MAX; NUM_PERMS];
299    for &x in shingles {
300        for (i, slot) in sig.iter_mut().enumerate() {
301            let h = splitmix64(x ^ SEEDS[i]);
302            if h < *slot {
303                *slot = h;
304            }
305        }
306    }
307    Some(sig)
308}
309
310/// Serialize a signature into a 1024-byte little-endian BLOB.
311pub fn signature_to_blob(sig: &[u64; NUM_PERMS]) -> Vec<u8> {
312    let mut blob = Vec::with_capacity(NUM_PERMS * 8);
313    for v in sig {
314        blob.extend_from_slice(&v.to_le_bytes());
315    }
316    blob
317}
318
319/// Deserialize a 1024-byte little-endian BLOB into a signature.
320pub fn blob_to_signature(blob: &[u8]) -> Option<[u64; NUM_PERMS]> {
321    if blob.len() != NUM_PERMS * 8 {
322        return None;
323    }
324    let mut sig = [0u64; NUM_PERMS];
325    for (i, chunk) in blob.chunks_exact(8).enumerate() {
326        sig[i] = u64::from_le_bytes(chunk.try_into().ok()?);
327    }
328    Some(sig)
329}
330
331/// Estimate Jaccard similarity from two MinHash signatures (fraction of
332/// matching permutation slots).
333pub fn estimate_similarity(a: &[u64; NUM_PERMS], b: &[u64; NUM_PERMS]) -> f64 {
334    let matching = a.iter().zip(b.iter()).filter(|(x, y)| x == y).count();
335    matching as f64 / NUM_PERMS as f64
336}
337
338/// Exact Jaccard similarity between two shingle sets.
339pub fn jaccard(a: &HashSet<u64>, b: &HashSet<u64>) -> f64 {
340    if a.is_empty() && b.is_empty() {
341        return 1.0;
342    }
343    if a.is_empty() || b.is_empty() {
344        return 0.0;
345    }
346    let intersection = a.intersection(b).count();
347    let union = a.len() + b.len() - intersection;
348    intersection as f64 / union as f64
349}
350
351/// Compute the LSH bucket key for each of the 16 bands (FNV-1a over the
352/// band's 8 signature words in little-endian byte order).
353pub fn band_keys(sig: &[u64; NUM_PERMS]) -> [u64; LSH_BANDS] {
354    let mut keys = [0u64; LSH_BANDS];
355    for (band, key) in keys.iter_mut().enumerate() {
356        let mut h = FNV_OFFSET;
357        for v in &sig[band * LSH_ROWS..(band + 1) * LSH_ROWS] {
358            for b in v.to_le_bytes() {
359                h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
360            }
361        }
362        *key = h;
363    }
364    keys
365}
366
367/// Compute fingerprints for all function/method symbols of one parsed file.
368///
369/// `symbols` are the parser-produced symbols (pre-id-rewrite); `id_map`
370/// translates their ids to the stored `path::[parent::]name@line` form.
371/// Symbols with fewer than [`SHINGLE_K`] tokens (no shingles) get no
372/// fingerprint. Returns an empty vec for languages [`tokenize`] cannot handle
373/// (e.g. YAML).
374pub fn file_fingerprints(
375    lang: Language,
376    source: &str,
377    rel_path: &str,
378    symbols: &[Symbol],
379    id_map: &HashMap<String, String>,
380) -> Vec<Fingerprint> {
381    let Some(tokens) = tokenize(lang, source) else {
382        return Vec::new();
383    };
384
385    let mut fingerprints = Vec::new();
386    for symbol in symbols {
387        if !matches!(symbol.kind, SymbolKind::Function | SymbolKind::Method) {
388            continue;
389        }
390        let Some(symbol_id) = id_map.get(&symbol.id) else {
391            continue;
392        };
393        let symbol_tokens: Vec<Tok> = tokens
394            .iter()
395            .filter(|t| t.line >= symbol.line_start && t.line <= symbol.line_end)
396            .cloned()
397            .collect();
398        let shingles = shingle_set(&symbol_tokens);
399        let Some(sig) = minhash(&shingles) else {
400            continue;
401        };
402        fingerprints.push(Fingerprint {
403            symbol_id: symbol_id.clone(),
404            file_path: rel_path.to_string(),
405            minhash: signature_to_blob(&sig),
406            token_count: symbol_tokens.len() as i64,
407        });
408    }
409    fingerprints
410}
411
412/// A verified near-duplicate pair, ordered so that `a.id < b.id`.
413#[derive(Debug, Clone)]
414pub struct DuplicatePair {
415    pub a: Symbol,
416    pub b: Symbol,
417    /// Exact Jaccard similarity over normalized token shingles.
418    pub similarity: f64,
419    pub token_count_a: i64,
420    pub token_count_b: i64,
421}
422
423/// Find near-duplicate function pairs in the index.
424///
425/// Loads all fingerprints with at least `min_tokens` tokens, buckets them
426/// with LSH banding, optionally keeps only pairs touching `changed_files`,
427/// then verifies each candidate pair with the exact Jaccard similarity over
428/// shingle sets re-derived from the symbols' stored source snippets.
429///
430/// Results are sorted by similarity (descending), then by the two symbol
431/// ids (ascending) for a stable order.
432pub fn find_near_duplicates(
433    db: &Database,
434    threshold: f64,
435    min_tokens: i64,
436    changed_files: Option<&HashSet<String>>,
437) -> Result<Vec<DuplicatePair>> {
438    let fingerprints = db.get_fingerprints(min_tokens)?;
439
440    // Decode signatures once (fingerprints are ordered by symbol_id, so
441    // index order is id order and (i, j) with i < j is the canonical pair).
442    let signatures: Vec<Option<[u64; NUM_PERMS]>> = fingerprints
443        .iter()
444        .map(|fp| blob_to_signature(&fp.minhash))
445        .collect();
446
447    // LSH banding: symbols sharing any band bucket become candidates.
448    let mut buckets: HashMap<(usize, u64), Vec<usize>> = HashMap::new();
449    for (idx, sig) in signatures.iter().enumerate() {
450        let Some(sig) = sig else { continue };
451        for (band, key) in band_keys(sig).iter().enumerate() {
452            buckets.entry((band, *key)).or_default().push(idx);
453        }
454    }
455
456    let mut candidates: HashSet<(usize, usize)> = HashSet::new();
457    for members in buckets.values() {
458        if members.len() < 2 {
459            continue;
460        }
461        for (n, &i) in members.iter().enumerate() {
462            for &j in &members[n + 1..] {
463                let pair = if i < j { (i, j) } else { (j, i) };
464                if pair.0 != pair.1 {
465                    candidates.insert(pair);
466                }
467            }
468        }
469    }
470
471    // --against filter: at least one endpoint must be in a changed file.
472    if let Some(changed) = changed_files {
473        candidates.retain(|&(i, j)| {
474            changed.contains(&fingerprints[i].file_path)
475                || changed.contains(&fingerprints[j].file_path)
476        });
477    }
478
479    // Verify candidates with exact Jaccard over re-derived shingle sets.
480    // Shingle sets are not stored; they are rebuilt by re-tokenizing each
481    // symbol's stored source snippet.
482    let mut shingle_cache: HashMap<usize, Option<HashSet<u64>>> = HashMap::new();
483    let mut symbol_cache: HashMap<usize, Option<Symbol>> = HashMap::new();
484    let mut pairs = Vec::new();
485
486    let mut sorted_candidates: Vec<(usize, usize)> = candidates.into_iter().collect();
487    sorted_candidates.sort_unstable();
488
489    for (i, j) in sorted_candidates {
490        for idx in [i, j] {
491            if let Entry::Vacant(entry) = symbol_cache.entry(idx) {
492                entry.insert(db.get_symbol(&fingerprints[idx].symbol_id)?);
493            }
494            if let Entry::Vacant(entry) = shingle_cache.entry(idx) {
495                entry.insert(symbol_cache[&idx].as_ref().and_then(symbol_shingles));
496            }
497        }
498
499        let (Some(sa), Some(sb)) = (&shingle_cache[&i], &shingle_cache[&j]) else {
500            continue;
501        };
502        let similarity = jaccard(sa, sb);
503        if similarity < threshold {
504            continue;
505        }
506        let (Some(a), Some(b)) = (&symbol_cache[&i], &symbol_cache[&j]) else {
507            continue;
508        };
509        pairs.push(DuplicatePair {
510            a: a.clone(),
511            b: b.clone(),
512            similarity,
513            token_count_a: fingerprints[i].token_count,
514            token_count_b: fingerprints[j].token_count,
515        });
516    }
517
518    pairs.sort_by(|x, y| {
519        y.similarity
520            .partial_cmp(&x.similarity)
521            .unwrap_or(std::cmp::Ordering::Equal)
522            .then_with(|| x.a.id.cmp(&y.a.id))
523            .then_with(|| x.b.id.cmp(&y.b.id))
524    });
525
526    Ok(pairs)
527}
528
529/// Rebuild a symbol's shingle set from its stored source snippet.
530///
531/// The language is detected from the file extension. Tree-sitter is
532/// error-tolerant, so tokenizing the snippet in isolation yields the same
533/// normalized token stream as slicing a whole-file parse (verified by test).
534/// Also used by `ctx score` to derive baseline shingles for symbols in
535/// unchanged files.
536pub fn symbol_shingles(symbol: &Symbol) -> Option<HashSet<u64>> {
537    let source = symbol.source.as_ref()?;
538    let lang = Language::from_path(Path::new(&symbol.file_path));
539    let tokens = tokenize(lang, source)?;
540    Some(shingle_set(&tokens))
541}
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546
547    fn texts(lang: Language, source: &str) -> Vec<String> {
548        tokenize(lang, source)
549            .expect("tokenize failed")
550            .into_iter()
551            .map(|t| t.text)
552            .collect()
553    }
554
555    #[test]
556    fn test_tokenize_rust_normalization() {
557        let src = "fn add(first: i32) -> i32 {\n    // a helpful remark\n    first + 42\n}\n";
558        let toks = texts(Language::Rust, src);
559
560        // Identifiers -> ID, numeric literals -> LIT, keywords kept.
561        assert!(toks.contains(&"ID".to_string()));
562        assert!(toks.contains(&"LIT".to_string()));
563        assert!(toks.contains(&"fn".to_string()));
564        // Raw identifier/literal text is gone.
565        assert!(!toks.contains(&"first".to_string()));
566        assert!(!toks.contains(&"42".to_string()));
567        // Comments are dropped entirely.
568        assert!(!toks.iter().any(|t| t.contains("remark")));
569
570        // Renamed identifiers + changed literals produce the same stream.
571        let renamed = "fn add(second: i32) -> i32 {\n    second + 7\n}\n";
572        assert_eq!(toks, texts(Language::Rust, renamed));
573    }
574
575    #[test]
576    fn test_tokenize_typescript_normalization() {
577        let src = "function greet(name: string): number {\n  /* block comment */\n  return name.length + 1;\n}\n";
578        let toks = texts(Language::TypeScript, src);
579        assert!(toks.contains(&"ID".to_string()));
580        assert!(toks.contains(&"LIT".to_string())); // 1 -> number -> LIT
581        assert!(toks.contains(&"function".to_string()));
582        assert!(!toks.contains(&"name".to_string()));
583        assert!(!toks.contains(&"1".to_string()));
584        assert!(!toks.iter().any(|t| t.contains("block comment")));
585    }
586
587    #[test]
588    fn test_tokenize_python_normalization() {
589        let src = "def add(x):\n    # trailing note\n    return x + 42\n";
590        let toks = texts(Language::Python, src);
591        assert!(toks.contains(&"ID".to_string()));
592        assert!(toks.contains(&"LIT".to_string())); // 42 -> integer -> LIT
593        assert!(toks.contains(&"def".to_string()));
594        assert!(!toks.contains(&"x".to_string()));
595        assert!(!toks.contains(&"42".to_string()));
596        assert!(!toks.iter().any(|t| t.contains("note")));
597    }
598
599    #[test]
600    fn test_tokenize_go_normalization() {
601        let src = "package main\n\nfunc add(a int) int {\n\t// short comment\n\treturn a + 42\n}\n";
602        let toks = texts(Language::Go, src);
603        assert!(toks.contains(&"ID".to_string()));
604        assert!(toks.contains(&"LIT".to_string())); // 42 -> int_literal -> LIT
605        assert!(toks.contains(&"func".to_string()));
606        assert!(!toks.contains(&"a".to_string()));
607        assert!(!toks.contains(&"42".to_string()));
608        assert!(!toks.iter().any(|t| t.contains("short comment")));
609    }
610
611    #[test]
612    fn test_tokenize_solidity_normalization() {
613        let src = "contract C {\n    function add(uint256 first) public pure returns (uint256) {\n        // a helpful remark\n        return first + 42;\n    }\n}\n";
614        let toks = texts(Language::Solidity, src);
615
616        // Identifiers -> ID, numeric literals -> LIT, keywords/punctuation kept.
617        assert!(toks.contains(&"ID".to_string()));
618        assert!(toks.contains(&"LIT".to_string()));
619        assert!(toks.contains(&"function".to_string()));
620        assert!(toks.contains(&"+".to_string()));
621        // Raw identifier/literal text is gone.
622        assert!(!toks.contains(&"first".to_string()));
623        assert!(!toks.contains(&"42".to_string()));
624        // Comments are dropped entirely (the solang lexer excludes them).
625        assert!(!toks.iter().any(|t| t.contains("remark")));
626
627        // Renamed identifiers + changed literals produce the same stream.
628        let renamed = "contract C {\n    function add(uint256 second) public pure returns (uint256) {\n        return second + 7;\n    }\n}\n";
629        assert_eq!(toks, texts(Language::Solidity, renamed));
630    }
631
632    #[test]
633    fn test_tokenize_unsupported_languages() {
634        // Languages with neither a tree-sitter grammar nor a lexer path.
635        assert!(tokenize(Language::Yaml, "a: 1").is_none());
636        assert!(tokenize(Language::Unknown, "whatever").is_none());
637    }
638
639    #[test]
640    fn test_shingle_set_requires_k_tokens() {
641        let short: Vec<Tok> = (0..SHINGLE_K - 1)
642            .map(|i| Tok {
643                text: format!("t{}", i),
644                line: 1,
645            })
646            .collect();
647        assert!(shingle_set(&short).is_empty());
648        assert!(minhash(&shingle_set(&short)).is_none());
649
650        let exact: Vec<Tok> = (0..SHINGLE_K)
651            .map(|i| Tok {
652                text: format!("t{}", i),
653                line: 1,
654            })
655            .collect();
656        assert_eq!(shingle_set(&exact).len(), 1);
657    }
658
659    #[test]
660    fn test_minhash_estimate_close_to_exact_jaccard() {
661        // Synthetic shingle sets with a known overlap.
662        let a: HashSet<u64> = (0..1000u64).map(splitmix64).collect();
663        let b: HashSet<u64> = (200..1200u64).map(splitmix64).collect();
664        let exact = jaccard(&a, &b);
665        assert!((exact - 800.0 / 1200.0).abs() < 1e-9);
666
667        let est = estimate_similarity(&minhash(&a).unwrap(), &minhash(&b).unwrap());
668        assert!(
669            (est - exact).abs() <= 0.08,
670            "estimate {} too far from exact {}",
671            est,
672            exact
673        );
674
675        // Identical and disjoint sets hit the extremes exactly.
676        assert_eq!(
677            estimate_similarity(&minhash(&a).unwrap(), &minhash(&a).unwrap()),
678            1.0
679        );
680        let c: HashSet<u64> = (10_000..10_500u64).map(splitmix64).collect();
681        let low = estimate_similarity(&minhash(&a).unwrap(), &minhash(&c).unwrap());
682        assert!(low <= 0.08, "disjoint sets estimated at {}", low);
683    }
684
685    #[test]
686    fn test_lsh_banding_finds_similar_pair() {
687        // Two shingle sets with Jaccard ~0.9: at least one of the 16 bands
688        // should collide (detection probability ~0.994 at J=0.85).
689        let a: HashSet<u64> = (0..200u64).map(splitmix64).collect();
690        let mut b = a.clone();
691        for i in 0..10u64 {
692            b.remove(&splitmix64(i));
693            b.insert(splitmix64(1_000_000 + i));
694        }
695        assert!(jaccard(&a, &b) > 0.89);
696
697        let keys_a = band_keys(&minhash(&a).unwrap());
698        let keys_b = band_keys(&minhash(&b).unwrap());
699        assert!(
700            keys_a.iter().zip(keys_b.iter()).any(|(x, y)| x == y),
701            "expected at least one shared LSH band for a ~0.9-similar pair"
702        );
703
704        // A dissimilar set should not collide in any band.
705        let c: HashSet<u64> = (5_000..5_200u64).map(splitmix64).collect();
706        let keys_c = band_keys(&minhash(&c).unwrap());
707        assert!(!keys_a.iter().zip(keys_c.iter()).any(|(x, y)| x == y));
708    }
709
710    #[test]
711    fn test_fingerprints_are_deterministic() {
712        let src = "fn f(a: i32) -> i32 {\n    let b = a * 2;\n    b + 1\n}\n";
713        let t1 = tokenize(Language::Rust, src).unwrap();
714        let t2 = tokenize(Language::Rust, src).unwrap();
715        assert_eq!(t1, t2);
716
717        let blob1 = signature_to_blob(&minhash(&shingle_set(&t1)).unwrap());
718        let blob2 = signature_to_blob(&minhash(&shingle_set(&t2)).unwrap());
719        assert_eq!(blob1, blob2);
720        assert_eq!(blob1.len(), NUM_PERMS * 8);
721
722        // Blob round-trip is lossless.
723        let sig = minhash(&shingle_set(&t1)).unwrap();
724        assert_eq!(blob_to_signature(&blob1).unwrap(), sig);
725        assert!(blob_to_signature(&blob1[..100]).is_none());
726    }
727
728    #[test]
729    fn test_snippet_tokenization_matches_whole_file_slice() {
730        // Verification re-tokenizes stored source snippets (e.g. a method
731        // without its impl block); tree-sitter's error tolerance must yield
732        // the same normalized stream as slicing a whole-file parse.
733        let file = "struct Foo {\n    x: i32,\n}\n\nimpl Foo {\n    fn double(&self) -> i32 {\n        let value = self.x * 2;\n        value + 1\n    }\n}\n";
734        let whole = tokenize(Language::Rust, file).unwrap();
735
736        // The method spans lines 6-9.
737        let sliced: Vec<String> = whole
738            .iter()
739            .filter(|t| t.line >= 6 && t.line <= 9)
740            .map(|t| t.text.clone())
741            .collect();
742        assert!(!sliced.is_empty());
743
744        let snippet: String = file.lines().skip(5).take(4).collect::<Vec<_>>().join("\n");
745        let snippet_toks: Vec<String> = tokenize(Language::Rust, &snippet)
746            .unwrap()
747            .into_iter()
748            .map(|t| t.text)
749            .collect();
750
751        assert_eq!(sliced, snippet_toks);
752    }
753
754    #[test]
755    fn test_jaccard_edge_cases() {
756        let empty: HashSet<u64> = HashSet::new();
757        let some: HashSet<u64> = [1, 2, 3].into_iter().collect();
758        assert_eq!(jaccard(&empty, &empty), 1.0);
759        assert_eq!(jaccard(&empty, &some), 0.0);
760        assert_eq!(jaccard(&some, &some), 1.0);
761    }
762}