1use 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
42pub const NUM_PERMS: usize = 128;
44
45pub const SHINGLE_K: usize = 5;
47
48pub const LSH_BANDS: usize = 16;
50
51pub const LSH_ROWS: usize = 8;
53
54pub const MIN_THRESHOLD: f64 = 0.5;
57
58const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
59const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
60
61const SEED_STATE: u64 = 0x0c7f_1d5e_a5ed_c0de;
63
64const 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
74const 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
88fn 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#[derive(Debug, Clone, PartialEq, Eq)]
99pub struct Tok {
100 pub text: String,
101 pub line: u32,
102}
103
104fn 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 static PARSERS: RefCell<HashMap<&'static str, Parser>> = RefCell::new(HashMap::new());
123}
124
125fn 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 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
154pub 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 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; }
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
212fn tokenize_solidity(source: &str) -> Vec<Tok> {
221 use solang_parser::lexer::{Lexer, Token};
222
223 let line_starts = line_start_offsets(source);
224
225 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
251fn 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
262fn offset_to_line(line_starts: &[usize], offset: usize) -> u32 {
265 line_starts.partition_point(|&s| s <= offset) as u32
266}
267
268pub 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
291pub 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
310pub 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
319pub 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
331pub 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
338pub 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
351pub 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
367pub 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#[derive(Debug, Clone)]
414pub struct DuplicatePair {
415 pub a: Symbol,
416 pub b: Symbol,
417 pub similarity: f64,
419 pub token_count_a: i64,
420 pub token_count_b: i64,
421}
422
423pub 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 let signatures: Vec<Option<[u64; NUM_PERMS]>> = fingerprints
443 .iter()
444 .map(|fp| blob_to_signature(&fp.minhash))
445 .collect();
446
447 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 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 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
529pub 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 assert!(toks.contains(&"ID".to_string()));
562 assert!(toks.contains(&"LIT".to_string()));
563 assert!(toks.contains(&"fn".to_string()));
564 assert!(!toks.contains(&"first".to_string()));
566 assert!(!toks.contains(&"42".to_string()));
567 assert!(!toks.iter().any(|t| t.contains("remark")));
569
570 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())); 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())); 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())); 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 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 assert!(!toks.contains(&"first".to_string()));
623 assert!(!toks.contains(&"42".to_string()));
624 assert!(!toks.iter().any(|t| t.contains("remark")));
626
627 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 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 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 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 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 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 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 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 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}