use crate::error::Result;
use crate::types::{Edge, Node, NodeId};
use std::collections::HashMap;
pub struct Deduplicator {
pub threshold: f64,
}
impl Deduplicator {
pub fn new(threshold: f64) -> Self {
Deduplicator { threshold }
}
pub fn default_threshold() -> f64 {
0.97
}
pub fn deduplicate(
&self,
nodes: Vec<Node>,
edges: Vec<Edge>,
) -> Result<(Vec<Node>, Vec<Edge>)> {
self.check_cross_project(&nodes)?;
let mut merged_nodes: HashMap<NodeId, Node> = HashMap::new();
let mut exact_index: HashMap<String, NodeId> = HashMap::new();
let mut id_mapping: HashMap<NodeId, NodeId> = HashMap::new();
for node in nodes {
let key = self.exact_match_key(&node);
if let Some(existing_id) = exact_index.get(&key) {
id_mapping.insert(node.id.clone(), existing_id.clone());
continue;
}
exact_index.insert(key, node.id.clone());
merged_nodes.insert(node.id.clone(), node);
}
let ids: Vec<NodeId> = merged_nodes.keys().cloned().collect();
for i in 0..ids.len() {
for j in (i + 1)..ids.len() {
let ni = merged_nodes[&ids[i]].clone();
let nj = merged_nodes[&ids[j]].clone();
if ni.source_file != nj.source_file {
continue;
}
let mut score = jaro_winkler(&ni.label, &nj.label);
if self.same_community(&ni, &nj) {
score = (score + 0.05).min(1.0);
}
if score > self.threshold && !self.is_variant_pair(&ni.label, &nj.label) {
let target_id = ids[j].clone();
let keep_id = ids[i].clone();
id_mapping.insert(target_id, keep_id);
merged_nodes.remove(&ids[j]);
}
}
}
let remapped_edges: Vec<Edge> = edges
.into_iter()
.map(|e| {
let new_src = id_mapping
.get(&e.source)
.cloned()
.unwrap_or(e.source.clone());
let new_tgt = id_mapping
.get(&e.target)
.cloned()
.unwrap_or(e.target.clone());
Edge {
source: new_src,
target: new_tgt,
..e
}
})
.collect();
let final_nodes: Vec<Node> = merged_nodes.into_values().collect();
Ok((final_nodes, remapped_edges))
}
fn check_cross_project(&self, nodes: &[Node]) -> Result<()> {
let repos: std::collections::HashSet<&str> = nodes
.iter()
.filter_map(|n| n.metadata.get("repo"))
.map(|s| s.as_str())
.collect();
if repos.len() > 1 {
return Err(crate::error::CodeSynapseError::Parse(
"cross-project dedup not supported: nodes from different repos found".to_string(),
));
}
Ok(())
}
fn same_community(&self, a: &Node, b: &Node) -> bool {
match (a.community, b.community) {
(Some(ca), Some(cb)) => ca == cb,
_ => false,
}
}
fn exact_match_key(&self, node: &Node) -> String {
format!("{}:{}", node.source_file, node.label)
}
fn is_variant_pair(&self, a: &str, b: &str) -> bool {
let (shorter, longer) = if a.len() <= b.len() { (a, b) } else { (b, a) };
longer.starts_with(shorter) && longer.len() > shorter.len()
}
pub fn llm_tiebreaker(&self, pairs: Vec<(Node, Node)>) -> Vec<(String, String, bool)> {
pairs
.into_iter()
.map(|(a, b)| {
let score = jaro_winkler(&a.label, &b.label);
if score >= self.threshold {
(a.id, b.id, true)
} else {
(a.id, b.id, false)
}
})
.collect()
}
}
impl Default for Deduplicator {
fn default() -> Self {
Deduplicator::new(Deduplicator::default_threshold())
}
}
fn jaro_winkler(a: &str, b: &str) -> f64 {
let a = a.to_lowercase();
let b = b.to_lowercase();
if a == b {
return 1.0;
}
let len_a = a.len();
let len_b = b.len();
if len_a == 0 || len_b == 0 {
return 0.0;
}
let max_dist = (std::cmp::max(len_a, len_b) / 2).saturating_sub(1);
let max_dist = std::cmp::max(max_dist, 0);
let mut a_matched = vec![false; len_a];
let mut b_matched = vec![false; len_b];
let mut matches = 0;
let mut transpositions = 0;
#[allow(clippy::needless_range_loop)]
for i in 0..len_a {
let start = i.saturating_sub(max_dist);
let end = std::cmp::min(len_b, i + max_dist + 1);
for j in start..end {
if b_matched[j] {
continue;
}
if a.as_bytes()[i] == b.as_bytes()[j] {
a_matched[i] = true;
b_matched[j] = true;
matches += 1;
break;
}
}
}
if matches == 0 {
return 0.0;
}
let mut k = 0;
#[allow(clippy::needless_range_loop)]
for i in 0..len_a {
if a_matched[i] {
while !b_matched[k] {
k += 1;
}
if a.as_bytes()[i] != b.as_bytes()[k] {
transpositions += 1;
}
k += 1;
}
}
let jaro = (matches as f64 / len_a as f64
+ matches as f64 / len_b as f64
+ (matches as f64 - transpositions as f64 / 2.0) / matches as f64)
/ 3.0;
let prefix_len = a
.chars()
.zip(b.chars())
.take_while(|(ca, cb)| ca == cb)
.count()
.min(4);
jaro + prefix_len as f64 * 0.1 * (1.0 - jaro)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn make_node(id: &str, label: &str, source_file: &str) -> Node {
Node {
id: id.to_string(),
label: label.to_string(),
file_type: "code".to_string(),
source_file: source_file.to_string(),
source_location: None,
community: None,
rationale: None,
docstring: None,
metadata: HashMap::new(),
}
}
#[test]
fn test_dedup_exact_same_file() {
let nodes = vec![
make_node("a", "AuthService", "auth.py"),
make_node("b", "AuthService", "auth.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 1);
}
#[test]
fn test_dedup_exact_diff_file() {
let nodes = vec![
make_node("a", "AuthService", "auth.py"),
make_node("b", "AuthService", "service.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 2);
}
#[test]
fn test_dedup_fuzzy_high_score() {
let nodes = vec![
make_node("a", "AuthService", "auth.py"),
make_node("b", "AuthService", "auth.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 1);
}
#[test]
fn test_dedup_fuzzy_low_score() {
let nodes = vec![
make_node("a", "AuthService", "auth.py"),
make_node("b", "AuthServiceImpl", "auth.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 2);
}
#[test]
fn test_dedup_variant_pair() {
let nodes = vec![
make_node("a", "M1", "config.py"),
make_node("b", "M1 Pro", "config.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 2);
}
#[test]
fn test_dedup_empty_nodes() {
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(vec![], vec![]).unwrap();
assert!(result.is_empty());
}
#[test]
fn test_dedup_remaps_edges() {
let nodes = vec![
make_node("a", "Service", "auth.py"),
make_node("a_dup", "Service", "auth.py"),
];
let edges = vec![Edge {
source: "a_dup".to_string(),
target: "b".to_string(),
relation: "calls".to_string(),
confidence: "EXTRACTED".to_string(),
source_file: Some("auth.py".to_string()),
weight: 1.0,
context: None,
}];
let dedup = Deduplicator::default();
let (result_nodes, result_edges) = dedup.deduplicate(nodes, edges).unwrap();
assert_eq!(result_nodes.len(), 1);
assert_eq!(result_edges[0].source, "a");
}
#[test]
fn test_jaro_winkler() {
let score = jaro_winkler("AuthService", "AuthService");
assert!((score - 1.0).abs() < 0.01);
let score = jaro_winkler("AuthService", "AuthServic");
assert!(score > 0.9);
}
#[test]
fn test_dedup_community_boost() {
let mut no = make_node("a", "abcXdefg", "auth.py");
let mut yes = make_node("b", "abcYdefh", "auth.py");
no.community = Some(1);
yes.community = Some(1);
let loose = Deduplicator::new(0.90);
let (r_loose, _) = loose
.deduplicate(vec![no.clone(), yes.clone()], vec![])
.unwrap();
assert_eq!(r_loose.len(), 1, "community boost pushes over 0.90");
let high = Deduplicator::new(0.97);
let (r_high, _) = high
.deduplicate(vec![no.clone(), yes.clone()], vec![])
.unwrap();
assert_eq!(r_high.len(), 2, "boosted score still below 0.97");
let no_community = Deduplicator::new(0.90);
let (r_nocom, _) = no_community
.deduplicate(
vec![
make_node("c", "abcXdefg", "auth.py"),
make_node("d", "abcYdefh", "auth.py"),
],
vec![],
)
.unwrap();
assert_eq!(r_nocom.len(), 2, "without boost, raw JW below 0.90");
}
#[test]
fn test_dedup_short_label() {
let nodes = vec![
make_node("a", "Extractor", "utils.py"),
make_node("b", "Extractar", "utils.py"),
];
let dedup = Deduplicator::new(0.92);
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(
result.len(),
1,
"short labels with 1-char diff should merge"
);
}
#[test]
fn test_dedup_short_label_blocked() {
let nodes = vec![
make_node("a", "cranel", "utils.py"),
make_node("b", "cranelr", "utils.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(
result.len(),
2,
"variant pair with different length blocked"
);
}
#[test]
fn test_dedup_cross_project_guard() {
let mut a = make_node("a", "Service", "a.py");
a.metadata.insert("repo".to_string(), "proj-a".to_string());
let mut b = make_node("b", "Service", "b.py");
b.metadata.insert("repo".to_string(), "proj-b".to_string());
let nodes = vec![a, b];
let dedup = Deduplicator::default();
let result = dedup.deduplicate(nodes, vec![]);
assert!(result.is_err(), "cross-project dedup should error");
}
#[test]
fn test_dedup_llm_tiebreaker() {
let a = make_node("a", "AuthService", "auth.py");
let b = make_node("b", "AuthServic", "auth.py");
let dedup = Deduplicator::default();
let decisions = dedup.llm_tiebreaker(vec![(a, b)]);
assert_eq!(decisions.len(), 1, "should return one decision");
}
#[test]
fn test_dedup_custom_threshold() {
let nodes = vec![
make_node("a", "AuthService", "auth.py"),
make_node("b", "AuthService", "auth.py"),
];
let loose = Deduplicator::new(0.99);
let (result, _) = loose.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 1);
let nodes = vec![
make_node("a", "AuthServid", "auth.py"),
make_node("b", "AuthServic", "auth.py"),
];
let strict = Deduplicator::new(0.99);
let (result_strict, _) = strict.deduplicate(nodes.clone(), vec![]).unwrap();
assert_eq!(result_strict.len(), 2, "high threshold blocks fuzzy merge");
let loose = Deduplicator::new(0.90);
let (result_loose, _) = loose.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result_loose.len(), 1, "low threshold allows fuzzy merge");
}
#[test]
fn test_dedup_all_unique() {
let nodes = vec![
make_node("a", "Alpha", "a.py"),
make_node("b", "Beta", "b.py"),
make_node("c", "Gamma", "c.py"),
];
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(nodes, vec![]).unwrap();
assert_eq!(result.len(), 3, "all unique labels → no dedup");
}
#[test]
fn test_dedup_edges_only_no_nodes() {
let edges = vec![Edge {
source: "a".to_string(),
target: "b".to_string(),
relation: "calls".to_string(),
confidence: "EXTRACTED".to_string(),
source_file: Some("test.py".to_string()),
weight: 1.0,
context: None,
}];
let dedup = Deduplicator::default();
let (result_nodes, result_edges) = dedup.deduplicate(vec![], edges).unwrap();
assert!(result_nodes.is_empty());
assert_eq!(result_edges.len(), 1, "edges pass through unchanged");
}
#[test]
fn test_dedup_default_threshold() {
let dedup = Deduplicator::default();
let (result, _) = dedup.deduplicate(vec![], vec![]).unwrap();
assert!(result.is_empty());
}
#[test]
fn test_jaro_winkler_identical() {
let score = jaro_winkler("hello", "hello");
assert!((score - 1.0).abs() < 0.001);
}
#[test]
fn test_jaro_winkler_empty_strings() {
let score = jaro_winkler("", "");
assert!(
(score - 1.0).abs() < 0.001,
"empty strings should return 1.0"
);
}
#[test]
fn test_jaro_winkler_completely_different() {
let score = jaro_winkler("abc", "xyz");
assert!(score < 0.5, "completely different should be low");
}
}