use std::collections::BTreeMap;
use crate::features::{FeatureHash, FeatureKind, FileFeatures};
pub const DEFAULT_POSTING_CAP: usize = 256;
pub const DEFAULT_PAIR_BUDGET: usize = 2_000_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidateConfig {
pub posting_cap: usize,
pub pair_budget: usize,
}
impl Default for CandidateConfig {
fn default() -> Self {
Self {
posting_cap: DEFAULT_POSTING_CAP,
pair_budget: DEFAULT_PAIR_BUDGET,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct StatementRun {
pub block: u32,
pub start: u32,
pub length: u32,
}
impl StatementRun {
#[must_use]
pub const fn end(self) -> u32 {
self.start.saturating_add(self.length)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct FragmentRef {
pub file: usize,
pub unit: usize,
pub start_byte: usize,
pub end_byte: usize,
pub extent: u32,
pub run: Option<StatementRun>,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct CandidatePair {
pub kind: FeatureKind,
pub hash: FeatureHash,
pub a: FragmentRef,
pub b: FragmentRef,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CandidateStats {
pub units: usize,
pub fragments: usize,
pub distinct_fingerprints: usize,
pub stop_fingerprints: usize,
pub stop_postings: usize,
pub candidate_pairs: usize,
pub available_pairs: usize,
pub budget_exhausted: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidateSet {
pub pairs: Vec<CandidatePair>,
pub stats: CandidateStats,
}
struct PairBudget {
remaining: usize,
exhausted: bool,
}
impl PairBudget {
const fn new(limit: usize) -> Self {
Self {
remaining: limit,
exhausted: false,
}
}
const fn take_list(&mut self, wanted: usize) -> bool {
if wanted > self.remaining {
self.exhausted = true;
return false;
}
self.remaining -= wanted;
true
}
}
#[must_use]
pub fn generate(files: &[FileFeatures], config: &CandidateConfig) -> CandidateSet {
let mut index: BTreeMap<(FeatureKind, FeatureHash), Vec<FragmentRef>> = BTreeMap::new();
let mut stats = CandidateStats::default();
for (file, features) in files.iter().enumerate() {
stats.units += features.units.len();
for (unit, unit_features) in features.units.iter().enumerate() {
for window in &unit_features.windows {
push_occurrence(
&mut index,
FeatureKind::StatementWindow,
window.hash,
FragmentRef {
file,
unit,
start_byte: window.range.start,
end_byte: window.range.end,
extent: clamp_u32(window.length),
run: Some(StatementRun {
block: window.block,
start: window.offset,
length: clamp_u32(window.length),
}),
},
);
stats.fragments += 1;
}
for subtree in &unit_features.subtrees {
push_occurrence(
&mut index,
FeatureKind::Subtree,
subtree.hash,
FragmentRef {
file,
unit,
start_byte: subtree.range.start,
end_byte: subtree.range.end,
extent: clamp_u32(subtree.node_count),
run: None,
},
);
stats.fragments += 1;
}
}
}
stats.distinct_fingerprints = index.len();
let mut eligible: Vec<(&(FeatureKind, FeatureHash), &Vec<FragmentRef>)> = Vec::new();
for (key, postings) in &index {
if postings.len() > config.posting_cap {
stats.stop_fingerprints += 1;
stats.stop_postings += postings.len();
} else if postings.len() >= 2 {
eligible.push((key, postings));
}
}
eligible.sort_by(|a, b| a.1.len().cmp(&b.1.len()).then_with(|| a.0.cmp(b.0)));
stats.available_pairs = eligible
.iter()
.map(|(_, postings)| pairs_within(postings.len()))
.sum();
let mut budget = PairBudget::new(config.pair_budget);
let mut pairs = Vec::new();
for (&(kind, hash), postings) in eligible {
if !budget.take_list(pairs_within(postings.len())) {
break;
}
for (i, &a) in postings.iter().enumerate() {
for &b in &postings[i + 1..] {
let (a, b) = if a <= b { (a, b) } else { (b, a) };
pairs.push(CandidatePair { kind, hash, a, b });
}
}
}
pairs.sort();
stats.candidate_pairs = pairs.len();
stats.budget_exhausted = budget.exhausted;
CandidateSet { pairs, stats }
}
fn push_occurrence(
index: &mut BTreeMap<(FeatureKind, FeatureHash), Vec<FragmentRef>>,
kind: FeatureKind,
hash: FeatureHash,
fragment: FragmentRef,
) {
index.entry((kind, hash)).or_default().push(fragment);
}
const fn pairs_within(len: usize) -> usize {
len.saturating_mul(len.saturating_sub(1)) / 2
}
fn clamp_u32(value: usize) -> u32 {
u32::try_from(value).unwrap_or(u32::MAX)
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests {
use super::*;
use crate::features::{
ApiCallFeature, CfgFeature, CharacteristicVector, SubtreeFeature, UnitFeatures,
WindowFeature,
};
use crate::ir::ByteRange;
fn hash(seed: u8) -> FeatureHash {
FeatureHash::from_bytes([seed; 16])
}
fn unit_with(windows: &[u8], subtrees: &[u8]) -> UnitFeatures {
let windows = windows
.iter()
.enumerate()
.map(|(i, &seed)| WindowFeature {
hash: hash(seed),
length: 4,
range: ByteRange {
start: i * 10,
end: i * 10 + 8,
},
block: 0,
offset: u32::try_from(i).unwrap(),
})
.collect();
let subtrees = subtrees
.iter()
.enumerate()
.map(|(i, &seed)| SubtreeFeature {
hash: hash(seed),
node_count: 6,
range: ByteRange {
start: 100 + i * 10,
end: 100 + i * 10 + 8,
},
})
.collect();
UnitFeatures {
name: None,
shape_tag: 1,
range: ByteRange { start: 0, end: 200 },
windows,
subtrees,
vector: CharacteristicVector::default(),
cfg: CfgFeature {
hash: hash(0),
skeleton_hash: hash(0),
op_count: 0,
skeleton_ops: 0,
max_loop_depth: 0,
branch_count: 0,
},
api: ApiCallFeature {
names: Vec::new(),
sequence_hash: hash(0),
multiset_hash: hash(0),
},
}
}
fn file_with(units: Vec<UnitFeatures>) -> FileFeatures {
FileFeatures { units }
}
#[test]
fn a_shared_hash_across_two_files_is_one_candidate_pair() {
let files = vec![
file_with(vec![unit_with(&[7], &[])]),
file_with(vec![unit_with(&[7], &[])]),
];
let set = generate(&files, &CandidateConfig::default());
assert_eq!(set.pairs.len(), 1);
let pair = &set.pairs[0];
assert_eq!(pair.kind, FeatureKind::StatementWindow);
assert_eq!(pair.hash, hash(7));
assert_eq!(pair.a.file, 0);
assert_eq!(pair.b.file, 1);
assert_eq!(set.stats.fragments, 2);
assert_eq!(set.stats.distinct_fingerprints, 1);
assert_eq!(set.stats.candidate_pairs, 1);
assert!(!set.stats.budget_exhausted);
}
#[test]
fn a_singleton_hash_yields_no_pair() {
let files = vec![file_with(vec![unit_with(&[7], &[8])])];
let set = generate(&files, &CandidateConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.distinct_fingerprints, 2);
assert_eq!(set.stats.stop_fingerprints, 0);
}
#[test]
fn window_and_subtree_hashes_do_not_cross_match() {
let files = vec![file_with(vec![unit_with(&[9], &[9])])];
let set = generate(&files, &CandidateConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.distinct_fingerprints, 2);
}
#[test]
fn a_high_frequency_hash_is_dropped_whole_and_counted() {
let files = vec![file_with(vec![
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
])];
let config = CandidateConfig {
posting_cap: 3,
..CandidateConfig::default()
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert_eq!(set.stats.stop_fingerprints, 1);
assert_eq!(set.stats.stop_postings, 4);
assert_eq!(set.stats.candidate_pairs, 0);
}
#[test]
fn the_pair_budget_refuses_a_list_it_cannot_pair_whole() {
let files = vec![file_with(vec![
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
])];
let config = CandidateConfig {
posting_cap: 64,
pair_budget: 2,
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert!(set.stats.budget_exhausted);
assert_eq!(set.stats.candidate_pairs, 0);
assert_eq!(set.stats.available_pairs, 6);
}
#[test]
fn a_budget_that_holds_one_list_and_not_the_next_pairs_the_first_whole() {
let files = vec![file_with(vec![
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[5], &[]),
unit_with(&[9], &[]),
unit_with(&[9], &[]),
])];
let config = CandidateConfig {
posting_cap: 64,
pair_budget: 5,
};
let set = generate(&files, &config);
assert_eq!(set.pairs.len(), 1);
assert_eq!(set.pairs[0].hash, hash(9));
assert!(set.stats.budget_exhausted);
assert_eq!(set.stats.available_pairs, 7);
}
#[test]
fn generation_is_deterministic() {
let files = vec![
file_with(vec![unit_with(&[7, 8], &[20])]),
file_with(vec![unit_with(&[8], &[20, 21])]),
];
let a = generate(&files, &CandidateConfig::default());
let b = generate(&files, &CandidateConfig::default());
assert_eq!(a, b);
assert_eq!(a.pairs.len(), 2);
}
}