use std::collections::{BTreeMap, BTreeSet};
use crate::clone_class::CloneClass;
use crate::verify::Confidence;
pub const GROUPING_VERSION: &str = "grouping-v1";
#[derive(Debug, Clone, PartialEq)]
pub struct GroupingConfig {
pub medoid_min_similarity: f64,
pub min_pairwise_similarity: f64,
pub sampling_threshold: usize,
pub sample_size: usize,
pub max_component: usize,
}
impl Default for GroupingConfig {
fn default() -> Self {
Self {
medoid_min_similarity: 0.60,
min_pairwise_similarity: 0.60,
sampling_threshold: 256,
sample_size: 32,
max_component: 1024,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SimilarityEdge {
pub a: usize,
pub b: usize,
pub similarity: f64,
pub breakdown: Option<crate::verify::SimilarityBreakdown>,
pub class: CloneClass,
pub confidence: Confidence,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GroupingUnit {
pub key: [u8; 16],
}
#[derive(Debug, Clone, PartialEq)]
pub struct StructuralGroup {
pub clone_type: CloneClass,
pub confidence: Confidence,
pub canonical: usize,
pub members: Vec<usize>,
pub medoid_similarities: Vec<f64>,
pub min_pairwise: f64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GroupingStats {
pub units: usize,
pub edges: usize,
pub components: usize,
pub oversized_components: usize,
pub groups: usize,
pub medoid_ejections: usize,
pub sampled_medoids: usize,
pub sampled_medoid_candidates: usize,
pub linkage_splits: usize,
pub singletons: usize,
}
#[derive(Debug, Clone, PartialEq)]
pub struct GroupingSet {
pub groups: Vec<StructuralGroup>,
piece_of: BTreeMap<usize, u32>,
pub stats: GroupingStats,
}
impl GroupingSet {
#[must_use]
pub fn severed_by_the_ceiling(&self, a: usize, b: usize) -> bool {
match (self.piece_of.get(&a), self.piece_of.get(&b)) {
(Some(left), Some(right)) => left != right,
_ => false,
}
}
}
#[must_use]
pub fn group(
units: &[GroupingUnit],
edges: &[SimilarityEdge],
config: &GroupingConfig,
) -> GroupingSet {
let sim = SimilarityGraph::build(units.len(), edges);
let mut stats = GroupingStats {
units: units.len(),
edges: edges.len(),
..GroupingStats::default()
};
let components = connected_components(units.len(), edges);
stats.components = components.len();
let mut groups = Vec::new();
let mut piece_of: BTreeMap<usize, u32> = BTreeMap::new();
let mut next_piece = 0u32;
for component in &components {
let cut = component.len() > piece_limit(config);
for piece in refinable_pieces(component, units, config, &mut stats) {
if cut {
for &member in &piece {
piece_of.insert(member, next_piece);
}
next_piece += 1;
}
refine_component(&piece, units, &sim, config, &mut groups, &mut stats);
}
}
groups.sort_by(|left, right| {
units[left.canonical]
.key
.cmp(&units[right.canonical].key)
.then(left.members.len().cmp(&right.members.len()))
.then_with(|| {
left.members
.iter()
.map(|&member| units[member].key)
.cmp(right.members.iter().map(|&member| units[member].key))
})
});
stats.groups = groups.len();
GroupingSet {
groups,
piece_of,
stats,
}
}
const fn piece_limit(config: &GroupingConfig) -> usize {
if config.max_component > 2 {
config.max_component
} else {
2
}
}
struct SimilarityGraph {
edges: BTreeMap<(usize, usize), EdgeData>,
}
#[derive(Debug, Clone, Copy)]
struct EdgeData {
similarity: f64,
class: CloneClass,
confidence: Confidence,
}
impl SimilarityGraph {
fn build(_unit_count: usize, edges: &[SimilarityEdge]) -> Self {
let mut map = BTreeMap::new();
for edge in edges {
if edge.a == edge.b {
continue;
}
let key = ordered(edge.a, edge.b);
let data = EdgeData {
similarity: edge.similarity,
class: edge.class,
confidence: edge.confidence,
};
map.entry(key)
.and_modify(|existing: &mut EdgeData| {
if edge.similarity > existing.similarity {
*existing = data;
}
})
.or_insert(data);
}
Self { edges: map }
}
fn similarity(&self, a: usize, b: usize) -> f64 {
if a == b {
return 1.0;
}
self.edges
.get(&ordered(a, b))
.map_or(0.0, |data| data.similarity)
}
fn edge(&self, a: usize, b: usize) -> Option<EdgeData> {
if a == b {
return None;
}
self.edges.get(&ordered(a, b)).copied()
}
}
const fn ordered(a: usize, b: usize) -> (usize, usize) {
if a <= b { (a, b) } else { (b, a) }
}
fn connected_components(unit_count: usize, edges: &[SimilarityEdge]) -> Vec<Vec<usize>> {
let mut parent: Vec<usize> = (0..unit_count).collect();
let mut connected = BTreeSet::new();
for edge in edges {
if edge.a != edge.b {
union(&mut parent, edge.a, edge.b);
connected.insert(edge.a);
connected.insert(edge.b);
}
}
let mut buckets: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
for node in connected {
let root = find(&mut parent, node);
buckets.entry(root).or_default().push(node);
}
buckets.into_values().collect()
}
fn find(parent: &mut [usize], node: usize) -> usize {
let mut root = node;
while parent[root] != root {
root = parent[root];
}
let mut current = node;
while parent[current] != root {
let next = parent[current];
parent[current] = root;
current = next;
}
root
}
fn union(parent: &mut [usize], a: usize, b: usize) {
let ra = find(parent, a);
let rb = find(parent, b);
if ra != rb {
if ra < rb {
parent[rb] = ra;
} else {
parent[ra] = rb;
}
}
}
fn refinable_pieces(
component: &[usize],
units: &[GroupingUnit],
config: &GroupingConfig,
stats: &mut GroupingStats,
) -> Vec<Vec<usize>> {
let limit = piece_limit(config);
if component.len() <= limit {
return vec![component.to_vec()];
}
stats.oversized_components += 1;
let mut ordered = component.to_vec();
ordered.sort_by_key(|&member| units[member].key);
let mut pieces = Vec::new();
let mut current = Vec::new();
let mut class_start = 0;
while class_start < ordered.len() {
let key = units[ordered[class_start]].key;
let class_end = ordered[class_start..]
.iter()
.position(|&member| units[member].key != key)
.map_or(ordered.len(), |offset| class_start + offset);
let class = &ordered[class_start..class_end];
if !current.is_empty() && current.len() + class.len() > limit {
pieces.push(std::mem::take(&mut current));
}
current.extend_from_slice(class);
if current.len() > limit {
pieces.push(std::mem::take(&mut current));
}
class_start = class_end;
}
if !current.is_empty() {
pieces.push(current);
}
pieces
}
fn refine_component(
component: &[usize],
units: &[GroupingUnit],
sim: &SimilarityGraph,
config: &GroupingConfig,
groups: &mut Vec<StructuralGroup>,
stats: &mut GroupingStats,
) {
if component.len() < 2 {
stats.singletons += component.len();
return;
}
let medoid = select_medoid(component, units, sim, config, stats);
let mut kept = Vec::new();
let mut rest = Vec::new();
for &member in component {
if member == medoid || sim.similarity(member, medoid) >= config.medoid_min_similarity {
kept.push(member);
} else {
rest.push(member);
}
}
stats.medoid_ejections += rest.len();
complete_linkage_trim(medoid, &mut kept, &mut rest, units, sim, config, stats);
if let Some(built) = build_group(medoid, &kept, units, sim) {
groups.push(built);
} else {
stats.singletons += kept.len();
}
if !rest.is_empty() {
rest.sort_by_key(|&m| units[m].key);
refine_component(&rest, units, sim, config, groups, stats);
}
}
fn select_medoid(
component: &[usize],
units: &[GroupingUnit],
sim: &SimilarityGraph,
config: &GroupingConfig,
stats: &mut GroupingStats,
) -> usize {
let mut candidates: Vec<usize> = component.to_vec();
candidates.sort_by_key(|&m| units[m].key);
if candidates.len() > config.sampling_threshold {
candidates.dedup_by_key(|member| units[*member].key);
let sample_size = config.sample_size.max(1).min(candidates.len());
if candidates.len() > sample_size {
let last = candidates.len() - 1;
candidates = if sample_size == 1 {
vec![candidates[last / 2]]
} else {
(0..sample_size)
.map(|index| candidates[index * last / (sample_size - 1)])
.collect()
};
}
stats.sampled_medoids += 1;
stats.sampled_medoid_candidates += candidates.len();
}
let mut best = candidates[0];
let mut best_total = total_similarity(best, component, sim);
for &candidate in &candidates[1..] {
let total = total_similarity(candidate, component, sim);
let better = match total.total_cmp(&best_total) {
std::cmp::Ordering::Greater => true,
std::cmp::Ordering::Equal => units[candidate].key < units[best].key,
std::cmp::Ordering::Less => false,
};
if better {
best = candidate;
best_total = total;
}
}
best
}
fn total_similarity(member: usize, set: &[usize], sim: &SimilarityGraph) -> f64 {
set.iter()
.filter(|&&other| other != member)
.map(|&other| sim.similarity(member, other))
.sum()
}
fn complete_linkage_trim(
medoid: usize,
kept: &mut Vec<usize>,
rest: &mut Vec<usize>,
units: &[GroupingUnit],
sim: &SimilarityGraph,
config: &GroupingConfig,
stats: &mut GroupingStats,
) {
let members = kept.clone();
let mut active = vec![true; members.len()];
let mut totals = vec![0.0; members.len()];
let mut pairs = Vec::with_capacity(kept.len().saturating_mul(kept.len().saturating_sub(1)) / 2);
for (index, &left) in members.iter().enumerate() {
for (right_index, &right) in members.iter().enumerate().skip(index + 1) {
let similarity = sim.similarity(left, right);
totals[index] += similarity;
totals[right_index] += similarity;
pairs.push((
similarity,
canonical_pair(left, right, units),
index,
right_index,
));
}
}
pairs.sort_by(|left, right| {
left.0
.total_cmp(&right.0)
.then_with(|| left.1.cmp(&right.1))
});
let mut next_pair = 0;
let mut active_count = members.len();
while active_count >= 2 {
while pairs
.get(next_pair)
.is_some_and(|(_, _, left, right)| !active[*left] || !active[*right])
{
next_pair += 1;
}
let Some(&(worst_sim, _, left, right)) = pairs.get(next_pair) else {
break;
};
if worst_sim >= config.min_pairwise_similarity {
break;
}
let victim = if members[left] == medoid {
right
} else if members[right] == medoid {
left
} else {
match totals[left].total_cmp(&totals[right]) {
std::cmp::Ordering::Less => left,
std::cmp::Ordering::Equal
if units[members[left]].key >= units[members[right]].key =>
{
left
}
std::cmp::Ordering::Greater | std::cmp::Ordering::Equal => right,
}
};
active[victim] = false;
for (index, &member) in members.iter().enumerate() {
if active[index] {
totals[index] -= sim.similarity(member, members[victim]);
}
}
active_count -= 1;
rest.push(members[victim]);
stats.linkage_splits += 1;
}
*kept = members
.into_iter()
.zip(active)
.filter_map(|(member, active)| active.then_some(member))
.collect();
}
fn canonical_pair(left: usize, right: usize, units: &[GroupingUnit]) -> ([u8; 16], [u8; 16]) {
let left_key = units[left].key;
let right_key = units[right].key;
if left_key <= right_key {
(left_key, right_key)
} else {
(right_key, left_key)
}
}
fn build_group(
medoid: usize,
kept: &[usize],
units: &[GroupingUnit],
sim: &SimilarityGraph,
) -> Option<StructuralGroup> {
if kept.len() < 2 {
return None;
}
let mut ordered_members: Vec<usize> = kept.iter().copied().filter(|&m| m != medoid).collect();
ordered_members.sort_by_key(|&m| units[m].key);
ordered_members.insert(0, medoid);
let medoid_similarities: Vec<f64> = ordered_members
.iter()
.map(|&member| sim.similarity(medoid, member))
.collect();
let mut clone_type = CloneClass::Type1;
let mut confidence = Confidence::High;
let mut min_pairwise = 1.0_f64;
for (i, &left) in ordered_members.iter().enumerate() {
for &right in &ordered_members[i + 1..] {
min_pairwise = min_pairwise.min(sim.similarity(left, right));
if let Some(data) = sim.edge(left, right) {
clone_type = weaker_class(clone_type, data.class);
confidence = weaker_confidence(confidence, data.confidence);
}
}
}
Some(StructuralGroup {
clone_type,
confidence,
canonical: medoid,
members: ordered_members,
medoid_similarities,
min_pairwise,
})
}
const fn weaker_class(a: CloneClass, b: CloneClass) -> CloneClass {
match (a, b) {
(CloneClass::Type3, _) | (_, CloneClass::Type3) => CloneClass::Type3,
(CloneClass::Type2, _) | (_, CloneClass::Type2) => CloneClass::Type2,
_ => CloneClass::Type1,
}
}
const fn weaker_confidence(a: Confidence, b: Confidence) -> Confidence {
match (a, b) {
(Confidence::Low, _) | (_, Confidence::Low) => Confidence::Low,
(Confidence::Medium, _) | (_, Confidence::Medium) => Confidence::Medium,
_ => Confidence::High,
}
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests;