use std::collections::BTreeMap;
use crate::features::{FeatureHash, FileFeatures, UnitRef};
pub const DEFAULT_MIN_OPS: u32 = 4;
pub const DEFAULT_MAX_LENGTH_RATIO: f64 = 3.0;
pub const DEFAULT_POSTING_CAP: usize = 256;
pub const DEFAULT_PAIR_BUDGET: usize = 2_000_000;
#[derive(Debug, Clone, PartialEq)]
pub struct ControlFlowConfig {
pub min_ops: u32,
pub max_length_ratio: f64,
pub posting_cap: usize,
pub pair_budget: usize,
}
impl Default for ControlFlowConfig {
fn default() -> Self {
Self {
min_ops: DEFAULT_MIN_OPS,
max_length_ratio: DEFAULT_MAX_LENGTH_RATIO,
posting_cap: DEFAULT_POSTING_CAP,
pair_budget: DEFAULT_PAIR_BUDGET,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct ControlFlowPair {
pub a: UnitRef,
pub b: UnitRef,
pub hash: FeatureHash,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ControlFlowStats {
pub units: usize,
pub indexed_units: usize,
pub skipped_shallow: usize,
pub distinct_skeletons: usize,
pub stop_skeletons: usize,
pub stop_postings: usize,
pub filtered_by_size: usize,
pub candidate_pairs: usize,
pub budget_exhausted: bool,
pub budget_dropped: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ControlFlowSet {
pub pairs: Vec<ControlFlowPair>,
pub stats: ControlFlowStats,
}
#[must_use]
pub fn generate(files: &[FileFeatures], config: &ControlFlowConfig) -> ControlFlowSet {
let mut index: BTreeMap<FeatureHash, Vec<UnitRef>> = BTreeMap::new();
let mut stats = ControlFlowStats::default();
for (file, features) in files.iter().enumerate() {
stats.units += features.units.len();
for (unit, unit_features) in features.units.iter().enumerate() {
if unit_features.cfg.skeleton_ops < config.min_ops {
stats.skipped_shallow += 1;
continue;
}
index
.entry(unit_features.cfg.skeleton_hash)
.or_default()
.push(UnitRef {
file,
unit,
node_count: unit_features.vector.node_count,
});
}
}
stats.indexed_units = index.values().map(Vec::len).sum();
stats.distinct_skeletons = index.len();
let mut eligible: Vec<(&FeatureHash, &Vec<UnitRef>)> = Vec::new();
for (hash, postings) in &index {
if postings.len() > config.posting_cap {
stats.stop_skeletons += 1;
stats.stop_postings += postings.len();
} else if postings.len() >= 2 {
eligible.push((hash, postings));
}
}
eligible.sort_by(|a, b| a.1.len().cmp(&b.1.len()).then_with(|| a.0.cmp(b.0)));
let mut pairs = Vec::new();
let mut remaining = config.pair_budget;
for (index, &(hash, postings)) in eligible.iter().enumerate() {
let possible = postings
.len()
.saturating_mul(postings.len().saturating_sub(1))
/ 2;
if possible > remaining {
stats.budget_exhausted = true;
stats.budget_dropped = eligible[index..].iter().fold(0, |total, (_, list)| {
total.saturating_add(list.len().saturating_mul(list.len().saturating_sub(1)) / 2)
});
break;
}
remaining -= possible;
let mut filtered = 0usize;
for (i, &a) in postings.iter().enumerate() {
for &b in &postings[i + 1..] {
if a.within_length_ratio(b, config.max_length_ratio) {
} else {
filtered += 1;
}
}
}
stats.filtered_by_size += filtered;
for (i, &a) in postings.iter().enumerate() {
for &b in &postings[i + 1..] {
if a.within_length_ratio(b, config.max_length_ratio) {
let (a, b) = if a <= b { (a, b) } else { (b, a) };
pairs.push(ControlFlowPair { a, b, hash: *hash });
}
}
}
}
pairs.sort();
stats.candidate_pairs = pairs.len();
ControlFlowSet { pairs, stats }
}
#[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(skeleton: u8, op_count: u32, node_count: u32) -> UnitFeatures {
UnitFeatures {
name: None,
shape_tag: 1,
range: ByteRange { start: 0, end: 100 },
windows: Vec::<WindowFeature>::new(),
subtrees: Vec::<SubtreeFeature>::new(),
vector: CharacteristicVector {
node_count,
..CharacteristicVector::default()
},
cfg: CfgFeature {
hash: hash(skeleton),
skeleton_hash: hash(skeleton),
op_count,
skeleton_ops: op_count,
max_loop_depth: 1,
branch_count: 1,
},
api: ApiCallFeature {
names: Vec::new(),
sequence_hash: hash(0),
multiset_hash: hash(0),
},
}
}
fn file(units: Vec<UnitFeatures>) -> FileFeatures {
FileFeatures { units }
}
#[test]
fn two_units_sharing_a_skeleton_are_a_candidate_pair() {
let files = vec![file(vec![unit(1, 4, 20)]), file(vec![unit(1, 4, 24)])];
let set = generate(&files, &ControlFlowConfig::default());
assert_eq!(set.pairs.len(), 1);
assert_eq!(set.pairs[0].hash, hash(1));
assert_eq!((set.pairs[0].a.file, set.pairs[0].b.file), (0, 1));
assert_eq!(set.stats.indexed_units, 2);
assert!(!set.stats.budget_exhausted);
}
#[test]
fn different_skeletons_do_not_pair() {
let files = vec![file(vec![unit(1, 4, 20), unit(2, 4, 20)])];
let set = generate(&files, &ControlFlowConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.distinct_skeletons, 2);
}
#[test]
fn a_skeleton_too_small_to_mean_anything_is_not_indexed() {
let files = vec![file(vec![unit(1, 3, 20), unit(1, 3, 20)])];
let set = generate(&files, &ControlFlowConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.skipped_shallow, 2);
assert_eq!(set.stats.indexed_units, 0);
}
#[test]
fn a_size_mismatched_pair_is_rejected_and_counted() {
let files = vec![file(vec![unit(1, 4, 10), unit(1, 4, 40)])];
let set = generate(&files, &ControlFlowConfig::default());
assert!(set.pairs.is_empty());
assert_eq!(set.stats.filtered_by_size, 1);
}
#[test]
fn a_common_skeleton_is_dropped_whole_and_counted() {
let files = vec![file(vec![
unit(1, 4, 20),
unit(1, 4, 20),
unit(1, 4, 20),
unit(1, 4, 20),
])];
let config = ControlFlowConfig {
posting_cap: 3,
..ControlFlowConfig::default()
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert_eq!(set.stats.stop_skeletons, 1);
assert_eq!(set.stats.stop_postings, 4);
}
#[test]
fn the_pair_budget_refuses_a_list_it_cannot_pair_whole() {
let files = vec![file(vec![
unit(1, 4, 20),
unit(1, 4, 20),
unit(1, 4, 20),
unit(1, 4, 20),
])];
let config = ControlFlowConfig {
pair_budget: 2,
..ControlFlowConfig::default()
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert!(set.stats.budget_exhausted);
assert_eq!(set.stats.budget_dropped, 6);
}
#[test]
fn a_refused_list_stops_before_later_quadratic_work() {
let files = vec![file(vec![
unit(1, 4, 20),
unit(1, 4, 20),
unit(1, 4, 20),
unit(2, 4, 20),
unit(2, 4, 20),
unit(2, 4, 100),
unit(2, 4, 400),
])];
let config = ControlFlowConfig {
pair_budget: 1,
..ControlFlowConfig::default()
};
let set = generate(&files, &config);
assert!(set.pairs.is_empty());
assert!(set.stats.budget_exhausted);
assert_eq!(set.stats.filtered_by_size, 0);
assert_eq!(set.stats.budget_dropped, 9);
}
#[test]
fn generation_is_deterministic() {
let files = vec![
file(vec![unit(1, 4, 20), unit(2, 5, 30)]),
file(vec![unit(1, 4, 22), unit(2, 5, 31)]),
];
let a = generate(&files, &ControlFlowConfig::default());
let b = generate(&files, &ControlFlowConfig::default());
assert_eq!(a, b);
assert_eq!(a.pairs.len(), 2);
}
}