use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use tree_sitter::Parser;
use crate::db::{Database, Fingerprint, Symbol, SymbolKind};
use crate::error::Result;
use crate::parser::Language;
pub const NUM_PERMS: usize = 128;
pub const SHINGLE_K: usize = 5;
pub const LSH_BANDS: usize = 16;
pub const LSH_ROWS: usize = 8;
pub const MIN_THRESHOLD: f64 = 0.5;
const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
const SEED_STATE: u64 = 0x0c7f_1d5e_a5ed_c0de;
const fn splitmix64(mut z: u64) -> u64 {
z = z.wrapping_add(0x9e37_79b9_7f4a_7c15);
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^ (z >> 31)
}
const SEEDS: [u64; NUM_PERMS] = {
let mut seeds = [0u64; NUM_PERMS];
let mut state = SEED_STATE;
let mut i = 0;
while i < NUM_PERMS {
state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
seeds[i] = splitmix64(state);
i += 1;
}
seeds
};
fn fnv1a(bytes: &[u8]) -> u64 {
let mut h = FNV_OFFSET;
for &b in bytes {
h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
}
h
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Tok {
pub text: String,
pub line: u32,
}
fn ts_language(lang: Language) -> Option<tree_sitter::Language> {
match lang {
Language::Rust => Some(tree_sitter_rust::language()),
Language::TypeScript => Some(tree_sitter_typescript::language_typescript()),
Language::Tsx => Some(tree_sitter_typescript::language_tsx()),
Language::JavaScript | Language::Jsx => Some(tree_sitter_javascript::language()),
Language::Python => Some(tree_sitter_python::language()),
Language::Go => Some(tree_sitter_go::language()),
Language::Solidity | Language::Yaml | Language::Unknown => None,
}
}
thread_local! {
static PARSERS: RefCell<HashMap<&'static str, Parser>> = RefCell::new(HashMap::new());
}
fn normalize_leaf(kind: &str, text: &str) -> String {
if kind.contains("identifier") {
return "ID".to_string();
}
if kind.contains("string")
|| kind.contains("char")
|| kind.contains("template")
|| kind.contains("rune")
{
return "LIT".to_string();
}
if matches!(
kind,
"integer_literal"
| "float_literal"
| "number"
| "int_literal"
| "imaginary_literal"
| "integer"
| "float"
) {
return "LIT".to_string();
}
text.to_string()
}
pub fn tokenize(lang: Language, source: &str) -> Option<Vec<Tok>> {
if lang == Language::Solidity {
return Some(tokenize_solidity(source));
}
let ts_lang = ts_language(lang)?;
PARSERS.with(|cell| {
let mut parsers = cell.borrow_mut();
let parser = parsers.entry(lang.as_str()).or_insert_with(|| {
let mut p = Parser::new();
p.set_language(ts_lang)
.expect("grammar/version mismatch for fingerprint parser");
p
});
let tree = parser.parse(source, None)?;
let bytes = source.as_bytes();
let mut tokens = Vec::new();
let mut stack = vec![tree.root_node()];
while let Some(node) = stack.pop() {
let kind = node.kind();
if kind.contains("comment") {
continue;
}
if node.child_count() == 0 {
let range = node.byte_range();
if range.is_empty() {
continue; }
let text = std::str::from_utf8(&bytes[range]).unwrap_or("");
tokens.push(Tok {
text: normalize_leaf(kind, text),
line: node.start_position().row as u32 + 1,
});
continue;
}
for i in (0..node.child_count()).rev() {
if let Some(child) = node.child(i) {
stack.push(child);
}
}
}
Some(tokens)
})
}
fn tokenize_solidity(source: &str) -> Vec<Tok> {
use solang_parser::lexer::{Lexer, Token};
let line_starts = line_start_offsets(source);
let mut comments = Vec::new();
let mut errors = Vec::new();
let lexer = Lexer::new(source, 0, &mut comments, &mut errors);
lexer
.map(|(start, token, _end)| {
let text = match token {
Token::Identifier(_) => "ID".to_string(),
Token::StringLiteral(..)
| Token::AddressLiteral(_)
| Token::HexLiteral(_)
| Token::Number(..)
| Token::RationalNumber(..)
| Token::HexNumber(_) => "LIT".to_string(),
other => other.to_string(),
};
Tok {
text,
line: offset_to_line(&line_starts, start),
}
})
.collect()
}
fn line_start_offsets(source: &str) -> Vec<usize> {
let mut starts = vec![0usize];
for (i, b) in source.bytes().enumerate() {
if b == b'\n' {
starts.push(i + 1);
}
}
starts
}
fn offset_to_line(line_starts: &[usize], offset: usize) -> u32 {
line_starts.partition_point(|&s| s <= offset) as u32
}
pub fn shingle_set(tokens: &[Tok]) -> HashSet<u64> {
let mut shingles = HashSet::new();
if tokens.len() < SHINGLE_K {
return shingles;
}
let hashes: Vec<u64> = tokens.iter().map(|t| fnv1a(t.text.as_bytes())).collect();
for window in hashes.windows(SHINGLE_K) {
let mut h = FNV_OFFSET;
for th in window {
for b in th.to_le_bytes() {
h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
}
}
shingles.insert(h);
}
shingles
}
pub fn minhash(shingles: &HashSet<u64>) -> Option<[u64; NUM_PERMS]> {
if shingles.is_empty() {
return None;
}
let mut sig = [u64::MAX; NUM_PERMS];
for &x in shingles {
for (i, slot) in sig.iter_mut().enumerate() {
let h = splitmix64(x ^ SEEDS[i]);
if h < *slot {
*slot = h;
}
}
}
Some(sig)
}
pub fn signature_to_blob(sig: &[u64; NUM_PERMS]) -> Vec<u8> {
let mut blob = Vec::with_capacity(NUM_PERMS * 8);
for v in sig {
blob.extend_from_slice(&v.to_le_bytes());
}
blob
}
pub fn blob_to_signature(blob: &[u8]) -> Option<[u64; NUM_PERMS]> {
if blob.len() != NUM_PERMS * 8 {
return None;
}
let mut sig = [0u64; NUM_PERMS];
for (i, chunk) in blob.chunks_exact(8).enumerate() {
sig[i] = u64::from_le_bytes(chunk.try_into().ok()?);
}
Some(sig)
}
pub fn estimate_similarity(a: &[u64; NUM_PERMS], b: &[u64; NUM_PERMS]) -> f64 {
let matching = a.iter().zip(b.iter()).filter(|(x, y)| x == y).count();
matching as f64 / NUM_PERMS as f64
}
pub fn jaccard(a: &HashSet<u64>, b: &HashSet<u64>) -> f64 {
if a.is_empty() && b.is_empty() {
return 1.0;
}
if a.is_empty() || b.is_empty() {
return 0.0;
}
let intersection = a.intersection(b).count();
let union = a.len() + b.len() - intersection;
intersection as f64 / union as f64
}
pub fn band_keys(sig: &[u64; NUM_PERMS]) -> [u64; LSH_BANDS] {
let mut keys = [0u64; LSH_BANDS];
for (band, key) in keys.iter_mut().enumerate() {
let mut h = FNV_OFFSET;
for v in &sig[band * LSH_ROWS..(band + 1) * LSH_ROWS] {
for b in v.to_le_bytes() {
h = (h ^ b as u64).wrapping_mul(FNV_PRIME);
}
}
*key = h;
}
keys
}
pub fn file_fingerprints(
lang: Language,
source: &str,
rel_path: &str,
symbols: &[Symbol],
id_map: &HashMap<String, String>,
) -> Vec<Fingerprint> {
let Some(tokens) = tokenize(lang, source) else {
return Vec::new();
};
let mut fingerprints = Vec::new();
for symbol in symbols {
if !matches!(symbol.kind, SymbolKind::Function | SymbolKind::Method) {
continue;
}
let Some(symbol_id) = id_map.get(&symbol.id) else {
continue;
};
let symbol_tokens: Vec<Tok> = tokens
.iter()
.filter(|t| t.line >= symbol.line_start && t.line <= symbol.line_end)
.cloned()
.collect();
let shingles = shingle_set(&symbol_tokens);
let Some(sig) = minhash(&shingles) else {
continue;
};
fingerprints.push(Fingerprint {
symbol_id: symbol_id.clone(),
file_path: rel_path.to_string(),
minhash: signature_to_blob(&sig),
token_count: symbol_tokens.len() as i64,
});
}
fingerprints
}
#[derive(Debug, Clone)]
pub struct DuplicatePair {
pub a: Symbol,
pub b: Symbol,
pub similarity: f64,
pub token_count_a: i64,
pub token_count_b: i64,
}
pub fn find_near_duplicates(
db: &Database,
threshold: f64,
min_tokens: i64,
changed_files: Option<&HashSet<String>>,
) -> Result<Vec<DuplicatePair>> {
let fingerprints = db.get_fingerprints(min_tokens)?;
let signatures: Vec<Option<[u64; NUM_PERMS]>> = fingerprints
.iter()
.map(|fp| blob_to_signature(&fp.minhash))
.collect();
let mut buckets: HashMap<(usize, u64), Vec<usize>> = HashMap::new();
for (idx, sig) in signatures.iter().enumerate() {
let Some(sig) = sig else { continue };
for (band, key) in band_keys(sig).iter().enumerate() {
buckets.entry((band, *key)).or_default().push(idx);
}
}
let mut candidates: HashSet<(usize, usize)> = HashSet::new();
for members in buckets.values() {
if members.len() < 2 {
continue;
}
for (n, &i) in members.iter().enumerate() {
for &j in &members[n + 1..] {
let pair = if i < j { (i, j) } else { (j, i) };
if pair.0 != pair.1 {
candidates.insert(pair);
}
}
}
}
if let Some(changed) = changed_files {
candidates.retain(|&(i, j)| {
changed.contains(&fingerprints[i].file_path)
|| changed.contains(&fingerprints[j].file_path)
});
}
let mut shingle_cache: HashMap<usize, Option<HashSet<u64>>> = HashMap::new();
let mut symbol_cache: HashMap<usize, Option<Symbol>> = HashMap::new();
let mut pairs = Vec::new();
let mut sorted_candidates: Vec<(usize, usize)> = candidates.into_iter().collect();
sorted_candidates.sort_unstable();
for (i, j) in sorted_candidates {
for idx in [i, j] {
if let Entry::Vacant(entry) = symbol_cache.entry(idx) {
entry.insert(db.get_symbol(&fingerprints[idx].symbol_id)?);
}
if let Entry::Vacant(entry) = shingle_cache.entry(idx) {
entry.insert(symbol_cache[&idx].as_ref().and_then(symbol_shingles));
}
}
let (Some(sa), Some(sb)) = (&shingle_cache[&i], &shingle_cache[&j]) else {
continue;
};
let similarity = jaccard(sa, sb);
if similarity < threshold {
continue;
}
let (Some(a), Some(b)) = (&symbol_cache[&i], &symbol_cache[&j]) else {
continue;
};
pairs.push(DuplicatePair {
a: a.clone(),
b: b.clone(),
similarity,
token_count_a: fingerprints[i].token_count,
token_count_b: fingerprints[j].token_count,
});
}
pairs.sort_by(|x, y| {
y.similarity
.partial_cmp(&x.similarity)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| x.a.id.cmp(&y.a.id))
.then_with(|| x.b.id.cmp(&y.b.id))
});
Ok(pairs)
}
pub fn symbol_shingles(symbol: &Symbol) -> Option<HashSet<u64>> {
let source = symbol.source.as_ref()?;
let lang = Language::from_path(Path::new(&symbol.file_path));
let tokens = tokenize(lang, source)?;
Some(shingle_set(&tokens))
}
#[cfg(test)]
mod tests {
use super::*;
fn texts(lang: Language, source: &str) -> Vec<String> {
tokenize(lang, source)
.expect("tokenize failed")
.into_iter()
.map(|t| t.text)
.collect()
}
#[test]
fn test_tokenize_rust_normalization() {
let src = "fn add(first: i32) -> i32 {\n // a helpful remark\n first + 42\n}\n";
let toks = texts(Language::Rust, src);
assert!(toks.contains(&"ID".to_string()));
assert!(toks.contains(&"LIT".to_string()));
assert!(toks.contains(&"fn".to_string()));
assert!(!toks.contains(&"first".to_string()));
assert!(!toks.contains(&"42".to_string()));
assert!(!toks.iter().any(|t| t.contains("remark")));
let renamed = "fn add(second: i32) -> i32 {\n second + 7\n}\n";
assert_eq!(toks, texts(Language::Rust, renamed));
}
#[test]
fn test_tokenize_typescript_normalization() {
let src = "function greet(name: string): number {\n /* block comment */\n return name.length + 1;\n}\n";
let toks = texts(Language::TypeScript, src);
assert!(toks.contains(&"ID".to_string()));
assert!(toks.contains(&"LIT".to_string())); assert!(toks.contains(&"function".to_string()));
assert!(!toks.contains(&"name".to_string()));
assert!(!toks.contains(&"1".to_string()));
assert!(!toks.iter().any(|t| t.contains("block comment")));
}
#[test]
fn test_tokenize_python_normalization() {
let src = "def add(x):\n # trailing note\n return x + 42\n";
let toks = texts(Language::Python, src);
assert!(toks.contains(&"ID".to_string()));
assert!(toks.contains(&"LIT".to_string())); assert!(toks.contains(&"def".to_string()));
assert!(!toks.contains(&"x".to_string()));
assert!(!toks.contains(&"42".to_string()));
assert!(!toks.iter().any(|t| t.contains("note")));
}
#[test]
fn test_tokenize_go_normalization() {
let src = "package main\n\nfunc add(a int) int {\n\t// short comment\n\treturn a + 42\n}\n";
let toks = texts(Language::Go, src);
assert!(toks.contains(&"ID".to_string()));
assert!(toks.contains(&"LIT".to_string())); assert!(toks.contains(&"func".to_string()));
assert!(!toks.contains(&"a".to_string()));
assert!(!toks.contains(&"42".to_string()));
assert!(!toks.iter().any(|t| t.contains("short comment")));
}
#[test]
fn test_tokenize_solidity_normalization() {
let src = "contract C {\n function add(uint256 first) public pure returns (uint256) {\n // a helpful remark\n return first + 42;\n }\n}\n";
let toks = texts(Language::Solidity, src);
assert!(toks.contains(&"ID".to_string()));
assert!(toks.contains(&"LIT".to_string()));
assert!(toks.contains(&"function".to_string()));
assert!(toks.contains(&"+".to_string()));
assert!(!toks.contains(&"first".to_string()));
assert!(!toks.contains(&"42".to_string()));
assert!(!toks.iter().any(|t| t.contains("remark")));
let renamed = "contract C {\n function add(uint256 second) public pure returns (uint256) {\n return second + 7;\n }\n}\n";
assert_eq!(toks, texts(Language::Solidity, renamed));
}
#[test]
fn test_tokenize_unsupported_languages() {
assert!(tokenize(Language::Yaml, "a: 1").is_none());
assert!(tokenize(Language::Unknown, "whatever").is_none());
}
#[test]
fn test_shingle_set_requires_k_tokens() {
let short: Vec<Tok> = (0..SHINGLE_K - 1)
.map(|i| Tok {
text: format!("t{}", i),
line: 1,
})
.collect();
assert!(shingle_set(&short).is_empty());
assert!(minhash(&shingle_set(&short)).is_none());
let exact: Vec<Tok> = (0..SHINGLE_K)
.map(|i| Tok {
text: format!("t{}", i),
line: 1,
})
.collect();
assert_eq!(shingle_set(&exact).len(), 1);
}
#[test]
fn test_minhash_estimate_close_to_exact_jaccard() {
let a: HashSet<u64> = (0..1000u64).map(splitmix64).collect();
let b: HashSet<u64> = (200..1200u64).map(splitmix64).collect();
let exact = jaccard(&a, &b);
assert!((exact - 800.0 / 1200.0).abs() < 1e-9);
let est = estimate_similarity(&minhash(&a).unwrap(), &minhash(&b).unwrap());
assert!(
(est - exact).abs() <= 0.08,
"estimate {} too far from exact {}",
est,
exact
);
assert_eq!(
estimate_similarity(&minhash(&a).unwrap(), &minhash(&a).unwrap()),
1.0
);
let c: HashSet<u64> = (10_000..10_500u64).map(splitmix64).collect();
let low = estimate_similarity(&minhash(&a).unwrap(), &minhash(&c).unwrap());
assert!(low <= 0.08, "disjoint sets estimated at {}", low);
}
#[test]
fn test_lsh_banding_finds_similar_pair() {
let a: HashSet<u64> = (0..200u64).map(splitmix64).collect();
let mut b = a.clone();
for i in 0..10u64 {
b.remove(&splitmix64(i));
b.insert(splitmix64(1_000_000 + i));
}
assert!(jaccard(&a, &b) > 0.89);
let keys_a = band_keys(&minhash(&a).unwrap());
let keys_b = band_keys(&minhash(&b).unwrap());
assert!(
keys_a.iter().zip(keys_b.iter()).any(|(x, y)| x == y),
"expected at least one shared LSH band for a ~0.9-similar pair"
);
let c: HashSet<u64> = (5_000..5_200u64).map(splitmix64).collect();
let keys_c = band_keys(&minhash(&c).unwrap());
assert!(!keys_a.iter().zip(keys_c.iter()).any(|(x, y)| x == y));
}
#[test]
fn test_fingerprints_are_deterministic() {
let src = "fn f(a: i32) -> i32 {\n let b = a * 2;\n b + 1\n}\n";
let t1 = tokenize(Language::Rust, src).unwrap();
let t2 = tokenize(Language::Rust, src).unwrap();
assert_eq!(t1, t2);
let blob1 = signature_to_blob(&minhash(&shingle_set(&t1)).unwrap());
let blob2 = signature_to_blob(&minhash(&shingle_set(&t2)).unwrap());
assert_eq!(blob1, blob2);
assert_eq!(blob1.len(), NUM_PERMS * 8);
let sig = minhash(&shingle_set(&t1)).unwrap();
assert_eq!(blob_to_signature(&blob1).unwrap(), sig);
assert!(blob_to_signature(&blob1[..100]).is_none());
}
#[test]
fn test_snippet_tokenization_matches_whole_file_slice() {
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";
let whole = tokenize(Language::Rust, file).unwrap();
let sliced: Vec<String> = whole
.iter()
.filter(|t| t.line >= 6 && t.line <= 9)
.map(|t| t.text.clone())
.collect();
assert!(!sliced.is_empty());
let snippet: String = file.lines().skip(5).take(4).collect::<Vec<_>>().join("\n");
let snippet_toks: Vec<String> = tokenize(Language::Rust, &snippet)
.unwrap()
.into_iter()
.map(|t| t.text)
.collect();
assert_eq!(sliced, snippet_toks);
}
#[test]
fn test_jaccard_edge_cases() {
let empty: HashSet<u64> = HashSet::new();
let some: HashSet<u64> = [1, 2, 3].into_iter().collect();
assert_eq!(jaccard(&empty, &empty), 1.0);
assert_eq!(jaccard(&empty, &some), 0.0);
assert_eq!(jaccard(&some, &some), 1.0);
}
}