use std::collections::{BTreeMap, BTreeSet};
use crate::features::{FileFeatures, UnitFeatures, UnitRef};
pub const DEFAULT_NUM_HASHES: usize = 128;
pub const DEFAULT_BANDS: usize = 64;
pub const DEFAULT_MAX_LENGTH_RATIO: f64 = 3.0;
pub const DEFAULT_MIN_SHINGLES: usize = 4;
pub const DEFAULT_MIN_ESTIMATED_JACCARD: f64 = 0.3;
pub const DEFAULT_POSTING_CAP: usize = 256;
pub const DEFAULT_PAIR_BUDGET: usize = 2_000_000;
pub const DEFAULT_MAX_SIGNED_UNITS: usize = 100_000;
pub const DEFAULT_NEAR_MISS_DELTA: f64 = 0.05;
pub const DEFAULT_NEAR_MISS_CAP: usize = 1_000;
#[derive(Debug, Clone, PartialEq)]
pub struct NearMatchConfig {
pub num_hashes: usize,
pub bands: usize,
pub min_shingles: usize,
pub max_length_ratio: f64,
pub min_estimated_jaccard: f64,
pub posting_cap: usize,
pub pair_budget: usize,
pub max_signed_units: usize,
pub near_miss_delta: f64,
pub near_miss_cap: usize,
}
impl Default for NearMatchConfig {
fn default() -> Self {
Self {
num_hashes: DEFAULT_NUM_HASHES,
bands: DEFAULT_BANDS,
min_shingles: DEFAULT_MIN_SHINGLES,
max_length_ratio: DEFAULT_MAX_LENGTH_RATIO,
min_estimated_jaccard: DEFAULT_MIN_ESTIMATED_JACCARD,
posting_cap: DEFAULT_POSTING_CAP,
pair_budget: DEFAULT_PAIR_BUDGET,
max_signed_units: DEFAULT_MAX_SIGNED_UNITS,
near_miss_delta: DEFAULT_NEAR_MISS_DELTA,
near_miss_cap: DEFAULT_NEAR_MISS_CAP,
}
}
}
impl NearMatchConfig {
fn rows(&self) -> usize {
(self.num_hashes / self.bands.max(1)).max(1)
}
fn near_miss_floor(&self) -> f64 {
(self.min_estimated_jaccard - self.near_miss_delta).max(0.0)
}
fn is_near_miss(&self, estimate: f64) -> bool {
estimate >= self.near_miss_floor() && estimate < self.min_estimated_jaccard
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NearMatchPair {
pub a: UnitRef,
pub b: UnitRef,
pub estimated_jaccard: f64,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NearMatchNearMiss {
pub a: UnitRef,
pub b: UnitRef,
pub estimated_jaccard: f64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct NearMatchStats {
pub units: usize,
pub signed_units: usize,
pub skipped_small: usize,
pub signed_limit_dropped: usize,
pub buckets: usize,
pub stop_buckets: usize,
pub stop_bucket_members: usize,
pub proposed_pairs: usize,
pub filtered_by_size: usize,
pub filtered_by_jaccard: usize,
pub near_miss_band_pairs: usize,
pub near_misses_retained: usize,
pub near_miss_cap_dropped: usize,
pub candidate_pairs: usize,
pub budget_exhausted: bool,
pub budget_dropped: usize,
}
#[derive(Debug, Clone, PartialEq)]
pub struct NearMatchSet {
pub pairs: Vec<NearMatchPair>,
pub near_misses: Vec<NearMatchNearMiss>,
pub stats: NearMatchStats,
}
#[must_use]
pub fn generate(files: &[FileFeatures], config: &NearMatchConfig) -> NearMatchSet {
let seeds = permutation_seeds(config.num_hashes);
let mut stats = NearMatchStats::default();
let mut signed = Vec::new();
let mut signatures = Vec::new();
for (file, features) in files.iter().enumerate() {
stats.units += features.units.len();
for (unit, unit_features) in features.units.iter().enumerate() {
let shingles = shingles_of(unit_features);
if shingles.len() < config.min_shingles {
stats.skipped_small += 1;
continue;
}
if signed.len() >= config.max_signed_units {
stats.signed_limit_dropped += 1;
continue;
}
let unit_ref = UnitRef {
file,
unit,
node_count: unit_features.vector.node_count,
};
signed.push(unit_ref);
signatures.extend(signature(&shingles, &seeds));
}
}
stats.signed_units = signed.len();
let proposed = propose_pairs(&signed, &signatures, config, &mut stats);
stats.proposed_pairs = proposed.len();
let mut pairs = Vec::new();
let mut near_misses = Vec::new();
for (ai, bi) in proposed {
let ref_a = signed[ai];
let ref_b = signed[bi];
if !ref_a.within_length_ratio(ref_b, config.max_length_ratio) {
stats.filtered_by_size += 1;
continue;
}
let estimated = estimated_jaccard(
signature_at(&signatures, ai, config.num_hashes),
signature_at(&signatures, bi, config.num_hashes),
);
if estimated < config.min_estimated_jaccard {
stats.filtered_by_jaccard += 1;
if config.is_near_miss(estimated) {
stats.near_miss_band_pairs += 1;
if near_misses.len() < config.near_miss_cap {
near_misses.push(NearMatchNearMiss {
a: ref_a,
b: ref_b,
estimated_jaccard: estimated,
});
} else {
stats.near_miss_cap_dropped += 1;
}
}
continue;
}
pairs.push(NearMatchPair {
a: ref_a,
b: ref_b,
estimated_jaccard: estimated,
});
}
stats.candidate_pairs = pairs.len();
stats.near_misses_retained = near_misses.len();
NearMatchSet {
pairs,
near_misses,
stats,
}
}
fn propose_pairs(
signed: &[UnitRef],
signatures: &[u64],
config: &NearMatchConfig,
stats: &mut NearMatchStats,
) -> Vec<(usize, usize)> {
let rows = config.rows();
let bands = config.num_hashes / rows;
let mut seen: BTreeSet<(usize, usize)> = BTreeSet::new();
let mut remaining = config.pair_budget;
for band in 0..bands {
let mut buckets: BTreeMap<u64, Vec<usize>> = BTreeMap::new();
for index in 0..signed.len() {
let signature = signature_at(signatures, index, config.num_hashes);
let start = band * rows;
let key = band_key(band, &signature[start..start + rows]);
buckets.entry(key).or_default().push(index);
}
let mut lists: Vec<Vec<usize>> = buckets
.into_values()
.filter(|members| members.len() >= 2)
.collect();
lists.sort();
lists.sort_by_key(Vec::len);
for members in lists {
stats.buckets += 1;
if members.len() > config.posting_cap {
stats.stop_buckets += 1;
stats.stop_bucket_members += members.len();
continue;
}
let mut unseen = Vec::new();
for (offset, &a) in members.iter().enumerate() {
for &b in &members[offset + 1..] {
let pair = if a <= b { (a, b) } else { (b, a) };
if !seen.contains(&pair) {
unseen.push(pair);
}
}
}
if unseen.len() > remaining {
stats.budget_exhausted = true;
stats.budget_dropped = unseen.len();
return seen.into_iter().collect();
}
remaining -= unseen.len();
seen.extend(unseen);
}
}
seen.into_iter().collect()
}
fn signature_at(signatures: &[u64], index: usize, width: usize) -> &[u64] {
let start = index.saturating_mul(width);
&signatures[start..start.saturating_add(width)]
}
fn shingles_of(unit: &UnitFeatures) -> Vec<u64> {
const WINDOW_DOMAIN: u64 = 0x5749_4e44_4f57_0000; const SUBTREE_DOMAIN: u64 = 0x5355_4254_5245_0000; let mut shingles: Vec<u64> = Vec::with_capacity(unit.windows.len() + unit.subtrees.len());
for window in &unit.windows {
shingles.push(fold_hash(window.hash.as_bytes()) ^ WINDOW_DOMAIN);
}
for subtree in &unit.subtrees {
shingles.push(fold_hash(subtree.hash.as_bytes()) ^ SUBTREE_DOMAIN);
}
shingles.sort_unstable();
shingles.dedup();
shingles
}
fn fold_hash(bytes: &[u8; 16]) -> u64 {
let mut lo = [0u8; 8];
let mut hi = [0u8; 8];
lo.copy_from_slice(&bytes[..8]);
hi.copy_from_slice(&bytes[8..]);
let a = u64::from_le_bytes(lo);
let b = u64::from_le_bytes(hi);
let mut z = a.wrapping_mul(0xff51_afd7_ed55_8ccd) ^ b.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
z = (z ^ (z >> 33)).wrapping_mul(0xff51_afd7_ed55_8ccd);
z ^ (z >> 29)
}
fn signature(shingles: &[u64], seeds: &[u64]) -> Vec<u64> {
seeds
.iter()
.map(|&seed| {
shingles
.iter()
.map(|&shingle| permute(shingle, seed))
.min()
.unwrap_or(u64::MAX)
})
.collect()
}
fn estimated_jaccard(a: &[u64], b: &[u64]) -> f64 {
let equal = a.iter().zip(b).filter(|(x, y)| x == y).count();
frac(equal, a.len())
}
fn frac(numer: usize, denom: usize) -> f64 {
let n = u32::try_from(numer).unwrap_or(u32::MAX);
let d = u32::try_from(denom).unwrap_or(u32::MAX);
if d == 0 {
0.0
} else {
f64::from(n) / f64::from(d)
}
}
fn permutation_seeds(count: usize) -> Vec<u64> {
let mut state = 0x1234_5678_9abc_def0u64;
(0..count).map(|_| splitmix64(&mut state)).collect()
}
const fn splitmix64(state: &mut u64) -> u64 {
*state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
let mut z = *state;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^ (z >> 31)
}
const fn permute(x: u64, seed: u64) -> u64 {
let mut z = x ^ seed;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^ (z >> 31)
}
fn band_key(band: usize, rows: &[u64]) -> u64 {
let mut z = 0xcbf2_9ce4_8422_2325u64 ^ (band as u64).wrapping_mul(0x1_0000_01b3);
for &row in rows {
z = (z ^ row).wrapping_mul(0x0000_0100_0000_01b3);
}
z ^ (z >> 32)
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests {
use super::*;
use crate::features::{
ApiCallFeature, CfgFeature, CharacteristicVector, FeatureHash, SubtreeFeature,
UnitFeatures, WindowFeature,
};
use crate::ir::ByteRange;
fn unit(windows: &[u8], subtrees: &[u8], node_count: u32) -> UnitFeatures {
let windows = windows
.iter()
.map(|&seed| WindowFeature {
hash: FeatureHash::from_bytes([seed; 16]),
length: 4,
range: ByteRange { start: 0, end: 8 },
block: 0,
offset: 0,
})
.collect();
let subtrees = subtrees
.iter()
.map(|&seed| SubtreeFeature {
hash: FeatureHash::from_bytes([seed; 16]),
node_count: 6,
range: ByteRange { start: 0, end: 8 },
})
.collect();
let vector = CharacteristicVector {
node_count,
..CharacteristicVector::default()
};
UnitFeatures {
name: None,
shape_tag: 1,
range: ByteRange { start: 0, end: 100 },
windows,
subtrees,
vector,
cfg: CfgFeature {
hash: FeatureHash::from_bytes([0; 16]),
skeleton_hash: FeatureHash::from_bytes([0; 16]),
op_count: 0,
skeleton_ops: 0,
max_loop_depth: 0,
branch_count: 0,
},
api: ApiCallFeature {
names: Vec::new(),
sequence_hash: FeatureHash::from_bytes([0; 16]),
multiset_hash: FeatureHash::from_bytes([0; 16]),
},
}
}
fn file(units: Vec<UnitFeatures>) -> FileFeatures {
FileFeatures { units }
}
#[test]
fn identical_units_are_a_candidate_with_full_similarity() {
let files = vec![
file(vec![unit(&[1, 2, 3, 4], &[5, 6], 20)]),
file(vec![unit(&[1, 2, 3, 4], &[5, 6], 20)]),
];
let set = generate(&files, &NearMatchConfig::default());
assert_eq!(set.pairs.len(), 1);
assert!((set.pairs[0].estimated_jaccard - 1.0).abs() < f64::EPSILON);
assert_eq!(set.stats.signed_units, 2);
assert!(!set.stats.budget_exhausted);
}
#[test]
fn signature_stage_stops_at_its_explicit_unit_ceiling() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
])];
let set = generate(
&files,
&NearMatchConfig {
max_signed_units: 2,
..NearMatchConfig::default()
},
);
assert_eq!(set.stats.signed_units, 2);
assert_eq!(set.stats.signed_limit_dropped, 1);
assert_eq!(set.pairs.len(), 1);
}
#[test]
fn a_high_overlap_pair_is_proposed_and_its_estimate_is_accurate() {
let a = unit(&[1, 2, 3, 4, 5], &[6, 7], 20);
let b = unit(&[1, 2, 3, 4, 5], &[8, 9], 20);
let files = vec![file(vec![a, b])];
let config = NearMatchConfig {
min_estimated_jaccard: 0.3,
..NearMatchConfig::default()
};
let set = generate(&files, &config);
assert_eq!(set.pairs.len(), 1, "a high-overlap pair must surface");
let true_jaccard = 5.0 / 9.0;
assert!(
(set.pairs[0].estimated_jaccard - true_jaccard).abs() < 0.15,
"estimate {} too far from {true_jaccard}",
set.pairs[0].estimated_jaccard
);
}
#[test]
fn disjoint_units_are_rejected_by_the_jaccard_gate() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[10, 11, 12, 13], &[14, 15], 20),
])];
let set = generate(&files, &NearMatchConfig::default());
assert!(
set.pairs.is_empty(),
"disjoint units must not be candidates"
);
assert_eq!(set.stats.candidate_pairs, 0);
}
#[test]
fn near_miss_band_includes_its_lower_bound_but_not_the_candidate_threshold() {
let config = NearMatchConfig {
min_estimated_jaccard: 0.75,
near_miss_delta: 0.25,
..NearMatchConfig::default()
};
assert!(config.is_near_miss(0.5));
assert!(config.is_near_miss(0.749_999));
assert!(!config.is_near_miss(0.499_999));
assert!(
!config.is_near_miss(0.75),
"an estimate that reaches the candidate threshold is never a near miss"
);
}
#[test]
fn near_miss_storage_is_capped_deterministically_without_changing_candidates() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
])];
let uncapped = NearMatchConfig {
min_estimated_jaccard: 1.1,
near_miss_delta: 1.1,
near_miss_cap: usize::MAX,
..NearMatchConfig::default()
};
let full = generate(&files, &uncapped);
let capped = NearMatchConfig {
near_miss_cap: 2,
..uncapped
};
let first = generate(&files, &capped);
let second = generate(&files, &capped);
assert!(full.pairs.is_empty());
assert_eq!(first.pairs, full.pairs);
assert_eq!(first.stats.candidate_pairs, full.stats.candidate_pairs);
assert_eq!(full.near_misses.len(), 3);
assert_eq!(first.near_misses.len(), 2);
assert_eq!(first.stats.near_miss_band_pairs, 3);
assert_eq!(first.stats.near_misses_retained, 2);
assert_eq!(first.stats.near_miss_cap_dropped, 1);
assert_eq!(first, second);
}
#[test]
fn the_length_ratio_gate_drops_size_mismatched_pairs() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 10),
unit(&[1, 2, 3, 4], &[5, 6], 40),
])];
let set = generate(&files, &NearMatchConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.filtered_by_size, 1);
assert_eq!(set.stats.filtered_by_jaccard, 0);
}
#[test]
fn a_unit_with_too_few_shingles_is_not_signed() {
let files = vec![file(vec![unit(&[1, 2], &[], 20), unit(&[1, 2], &[], 20)])];
let set = generate(&files, &NearMatchConfig::default());
assert_eq!(set.stats.signed_units, 0);
assert_eq!(set.stats.skipped_small, 2);
assert!(set.pairs.is_empty());
}
#[test]
fn a_high_frequency_bucket_is_dropped_and_counted() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
])];
let config = NearMatchConfig {
posting_cap: 3,
..NearMatchConfig::default()
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert!(set.stats.stop_buckets > 0);
assert_eq!(set.stats.candidate_pairs, 0);
}
#[test]
fn pair_budget_charges_each_distinct_pair_once_across_lsh_bands() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
])];
let set = generate(
&files,
&NearMatchConfig {
pair_budget: 3,
..NearMatchConfig::default()
},
);
assert_eq!(set.stats.proposed_pairs, 3);
assert!(!set.stats.budget_exhausted);
}
#[test]
fn the_pair_budget_refuses_a_bucket_it_cannot_hold_whole() {
let files = vec![file(vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
])];
let config = NearMatchConfig {
pair_budget: 1,
..NearMatchConfig::default()
};
let set = generate(&files, &config);
assert_eq!(set.stats.proposed_pairs, 0);
assert!(set.stats.budget_exhausted);
assert_eq!(set.stats.budget_dropped, 3);
}
#[test]
fn a_refused_bucket_stops_before_quadratic_deduplication() {
let units = vec![
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[1, 2, 3, 4], &[5, 6], 20),
unit(&[40, 41, 42, 43], &[44, 45], 20),
unit(&[40, 41, 42, 43], &[44, 45], 20),
];
let files = vec![file(units)];
let full = generate(&files, &NearMatchConfig::default());
let squeezed = generate(
&files,
&NearMatchConfig {
pair_budget: 1,
..NearMatchConfig::default()
},
);
assert!(squeezed.stats.budget_exhausted);
assert_eq!(squeezed.stats.proposed_pairs, 1);
assert!(
squeezed.stats.buckets < full.stats.buckets,
"the ceiling stops before walking buckets it cannot examine"
);
}
#[test]
fn generation_is_deterministic() {
let files = vec![
file(vec![unit(&[1, 2, 3, 4], &[5, 6], 20)]),
file(vec![unit(&[1, 2, 3, 5], &[5, 6], 22)]),
];
let a = generate(&files, &NearMatchConfig::default());
let b = generate(&files, &NearMatchConfig::default());
assert_eq!(a, b);
}
}