use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct BoundaryProposal {
pub community_id: u32,
pub stable_snapshots: u32,
pub node_ids: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct BoundaryInferenceResult {
pub proposals: Vec<BoundaryProposal>,
pub partition_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PriorPartition {
pub cluster_assignments: BTreeMap<String, u32>,
}
pub fn infer_boundaries(
partitions: &[PriorPartition],
stability_threshold: u32,
) -> BoundaryInferenceResult {
let partition_count = partitions.len();
if partition_count == 0 || stability_threshold == 0 {
return BoundaryInferenceResult {
proposals: vec![],
partition_count,
};
}
let latest = partitions.last().unwrap();
let latest_communities = invert_assignments(&latest.cluster_assignments);
let mut proposals: Vec<BoundaryProposal> = Vec::new();
for (comm_id, node_set) in &latest_communities {
let stable = count_stable_tail(partitions, node_set);
if stable >= stability_threshold {
let mut node_ids: Vec<String> = node_set.iter().cloned().collect();
node_ids.sort();
proposals.push(BoundaryProposal {
community_id: *comm_id,
stable_snapshots: stable,
node_ids,
});
}
}
proposals.sort_by_key(|p| p.community_id);
BoundaryInferenceResult {
proposals,
partition_count,
}
}
fn invert_assignments(
assignments: &BTreeMap<String, u32>,
) -> BTreeMap<u32, std::collections::BTreeSet<String>> {
let mut result: BTreeMap<u32, std::collections::BTreeSet<String>> = BTreeMap::new();
for (node, &comm) in assignments {
result.entry(comm).or_default().insert(node.clone());
}
result
}
fn count_stable_tail(
partitions: &[PriorPartition],
node_set: &std::collections::BTreeSet<String>,
) -> u32 {
if partitions.len() < 2 {
return 0;
}
let mut stable = 0u32;
for i in (0..partitions.len() - 1).rev() {
let p = &partitions[i];
let communities = invert_assignments(&p.cluster_assignments);
let found = communities.values().any(|s| s == node_set);
if found {
stable += 1;
} else {
break;
}
}
stable
}
#[cfg(test)]
mod tests {
use super::*;
fn partition_from(pairs: &[(&str, u32)]) -> PriorPartition {
PriorPartition {
cluster_assignments: pairs.iter().map(|(k, v)| (k.to_string(), *v)).collect(),
}
}
#[test]
fn empty_input_returns_no_proposals() {
let r = infer_boundaries(&[], 3);
assert!(r.proposals.is_empty());
assert_eq!(r.partition_count, 0);
}
#[test]
fn single_partition_no_proposals_due_to_no_history() {
let p = partition_from(&[("a", 0), ("b", 0), ("c", 1)]);
let r = infer_boundaries(&[p], 2);
assert!(r.proposals.is_empty());
}
#[test]
fn stable_community_proposed() {
let p1 = partition_from(&[("a", 0), ("b", 0), ("c", 1)]);
let p2 = partition_from(&[("a", 0), ("b", 0), ("c", 1)]);
let p3 = partition_from(&[("a", 0), ("b", 0), ("c", 1)]);
let r = infer_boundaries(&[p1, p2, p3], 2);
assert!(!r.proposals.is_empty());
let prop = r.proposals.iter().find(|p| p.node_ids == vec!["a", "b"]);
assert!(prop.is_some());
}
#[test]
fn threshold_not_met_no_proposal() {
let p1 = partition_from(&[("a", 0), ("b", 0)]);
let p2 = partition_from(&[("a", 0), ("b", 0)]);
let r = infer_boundaries(&[p1, p2], 3);
assert!(r.proposals.is_empty());
}
}