1use std::collections::HashMap;
15
16use petgraph::visit::{EdgeRef, IntoEdgeReferences, IntoNodeReferences};
17
18use leiden_rs::{GraphDataBuilder, Leiden, LeidenConfig, QualityType};
19
20use crate::model::*;
21
22const LEIDEN_SEED: u64 = 42;
24const LEIDEN_RESOLUTION: f64 = 0.5;
29
30pub const WEIGHT_IMPORTS: f64 = 0.8;
32pub const WEIGHT_CALLS: f64 = 0.7;
33
34fn file_dir_key(graph: &KnowledgeGraph, nid: NodeId) -> String {
36 let path = graph
37 .graph
38 .node_weight(nid)
39 .and_then(|n| n.file_path.as_deref())
40 .unwrap_or("");
41 let norm = path.replace('\\', "/");
42 let dir = norm.rsplit_once('/').map(|(d, _)| d).unwrap_or(".");
43 let dir = dir.trim_end_matches('/');
44 if dir.is_empty() {
45 "<root>".to_string()
46 } else {
47 dir.to_string()
48 }
49}
50
51pub fn detect_communities(graph: &KnowledgeGraph) -> Vec<Vec<NodeId>> {
55 detect_communities_with_resolution(graph, LEIDEN_RESOLUTION)
56}
57
58pub fn detect_communities_with_resolution(graph: &KnowledgeGraph, resolution: f64) -> Vec<Vec<NodeId>> {
61 let file_nodes: Vec<NodeId> = graph
62 .graph
63 .node_references()
64 .filter(|(_, n)| n.kind == NodeKind::File)
65 .map(|(id, _)| id)
66 .collect();
67
68 if file_nodes.is_empty() {
69 return Vec::new();
70 }
71 if file_nodes.len() == 1 {
72 return vec![file_nodes];
73 }
74
75 const MIN_DIRS_FOR_SUPERNODE: usize = 24;
91 let mut dirs: std::collections::BTreeMap<String, Vec<NodeId>> = std::collections::BTreeMap::new();
92 for &nid in &file_nodes {
93 let d = file_dir_key(graph, nid);
94 dirs.entry(d).or_default().push(nid);
95 }
96 if dirs.len() <= 1 || dirs.len() >= MIN_DIRS_FOR_SUPERNODE {
102 return dirs.into_values().collect();
103 }
104
105 let compact: HashMap<NodeId, usize> = file_nodes
107 .iter()
108 .enumerate()
109 .map(|(i, &nid)| (nid, i))
110 .collect();
111
112 let mut entity_to_file: HashMap<NodeId, NodeId> = HashMap::new();
116 for edge in graph.graph.edge_references() {
117 let kind = graph.graph.edge_weight(edge.id()).map(|e| e.kind.clone());
118 if kind == Some(EdgeKind::Contains) && compact.contains_key(&edge.source()) {
119 entity_to_file.insert(edge.target(), edge.source());
120 }
121 }
122 let file_of = |nid: NodeId| -> Option<NodeId> {
124 if compact.contains_key(&nid) {
125 Some(nid)
126 } else {
127 entity_to_file.get(&nid).copied()
128 }
129 };
130
131 let mut edge_weights: HashMap<(usize, usize), f64> = HashMap::new();
133 for edge in graph.graph.edge_references() {
134 let e = graph
138 .graph
139 .edge_weight(edge.id())
140 .expect("边权重必然存在");
141 let w = match e.kind {
142 EdgeKind::Imports => WEIGHT_IMPORTS,
143 EdgeKind::Calls => WEIGHT_CALLS,
144 _ => continue, };
146 let (Some(sf), Some(tf)) = (file_of(edge.source()), file_of(edge.target())) else {
147 continue; };
149 let (Some(&si), Some(&ti)) = (compact.get(&sf), compact.get(&tf)) else {
150 continue;
151 };
152 if si == ti {
153 continue; }
155 *edge_weights.entry((si, ti)).or_insert(0.0) += w;
156 }
157
158 if edge_weights.is_empty() {
159 return file_nodes.into_iter().map(|nid| vec![nid]).collect();
161 }
162
163 let mut builder = GraphDataBuilder::new(file_nodes.len()).directed();
168 for ((s, t), w) in &edge_weights {
169 builder
170 .add_edge(*s, *t, *w)
171 .expect("边权重均为有限非负数");
172 }
173 let data = builder.build().expect("图数据构造失败");
174
175 let config = LeidenConfig {
176 quality: QualityType::CPM,
177 resolution,
178 seed: Some(LEIDEN_SEED),
179 ..Default::default()
180 };
181 let result = Leiden::new(config)
186 .run(&data)
187 .expect("Leiden 社区检测失败");
188 let membership = result.partition.as_slice(); let mut groups: HashMap<usize, Vec<NodeId>> = HashMap::new();
192 for (i, &comm) in membership.iter().enumerate() {
193 groups.entry(comm).or_default().push(file_nodes[i]);
194 }
195
196 let mut communities: Vec<Vec<NodeId>> = groups.into_values().collect();
200 communities.sort_by(|a, b| {
201 b.len()
202 .cmp(&a.len())
203 .then_with(|| min_file_path(graph, a).cmp(&min_file_path(graph, b)))
204 });
205 communities
206}
207
208fn min_file_path(graph: &KnowledgeGraph, files: &[NodeId]) -> String {
210 files
211 .iter()
212 .filter_map(|nid| graph.graph.node_weight(*nid).and_then(|n| n.file_path.clone()))
213 .min()
214 .unwrap_or_default()
215}
216
217pub fn community_name(files: &[String], fallback_index: usize) -> String {
227 if files.is_empty() {
228 return format!("module_{fallback_index}");
229 }
230
231 let dirs: Vec<Vec<String>> = files.iter().map(|p| dir_segments(p)).collect();
232
233 let min_len = dirs.iter().map(|d| d.len()).min().unwrap_or(0);
235 let mut common = 0usize;
236 'outer: for i in 0..min_len {
237 let seg = &dirs[0][i];
238 for other in &dirs[1..] {
239 if other.get(i) != Some(seg) {
240 break 'outer;
241 }
242 }
243 common = i + 1;
244 }
245 if common > 0 {
246 return dirs[0][..common].join("::");
247 }
248
249 let mut dir_counts: HashMap<String, usize> = HashMap::new();
251 for d in &dirs {
252 let key = if d.is_empty() {
253 "<root>".to_string()
254 } else {
255 d.join("::")
256 };
257 *dir_counts.entry(key).or_insert(0) += 1;
258 }
259 if let Some((best, _)) = dir_counts
260 .into_iter()
261 .max_by_key(|(name, count)| (*count, name.clone()))
262 && best != "<root>"
263 {
264 return best;
265 }
266
267 format!("module_{fallback_index}")
269}
270
271fn dir_segments(path: &str) -> Vec<String> {
273 use std::path::Component;
274 std::path::Path::new(path)
275 .parent()
276 .map(|p| {
277 p.components()
278 .filter_map(|c| match c {
279 Component::Normal(s) => Some(s.to_string_lossy().into_owned()),
280 _ => None,
281 })
282 .collect()
283 })
284 .unwrap_or_default()
285}
286
287#[cfg(test)]
288mod tests {
289 use super::*;
290
291 fn make_graph() -> KnowledgeGraph {
293 let mut kg = KnowledgeGraph::default();
294 let g = &mut kg.graph;
295 let add_file =
296 |g: &mut petgraph::stable_graph::StableDiGraph<CodeNode, CodeEdge>,
297 path: &str|
298 -> (NodeId, NodeId) {
299 let module_path: Vec<String> = dir_segments(path);
300 let nid = g.add_node(CodeNode {
301 id: NodeId::new(g.node_count()),
302 kind: NodeKind::File,
303 name: path.into(),
304 file_path: Some(path.into()),
305 line_range: None,
306 doc_comment: None,
307 signature: None, visibility: None,
308 module_path,
309 });
310 let eid = g.add_node(CodeNode {
312 id: NodeId::new(g.node_count()),
313 kind: NodeKind::Function,
314 name: format!("f{}", nid.index()),
315 file_path: Some(path.into()),
316 line_range: None,
317 doc_comment: None,
318 signature: None, visibility: None,
319 module_path: Vec::new(),
320 });
321 g.add_edge(
322 nid,
323 eid,
324 CodeEdge {
325 id: EdgeId::new(g.edge_count()),
326 kind: EdgeKind::Contains,
327 source: nid,
328 target: eid,
329 weight: 1.0,
330 location: None,
331 },
332 );
333 (nid, eid)
334 };
335 let (_a, ea) = add_file(g, "src/a.rs");
336 let (_b, eb) = add_file(g, "src/b.rs");
337 let _tcp = add_file(g, "src/net/tcp.rs");
338 g.add_edge(
340 ea,
341 eb,
342 CodeEdge {
343 id: EdgeId::new(g.edge_count()),
344 kind: EdgeKind::Calls,
345 source: ea,
346 target: eb,
347 weight: 0.7,
348 location: None,
349 },
350 );
351 kg
352 }
353
354 #[test]
355 fn test_detect_communities_basic() {
356 let kg = make_graph();
357 let communities = detect_communities(&kg);
358 assert_eq!(communities.len(), 2, "应产出 2 个社区");
360 let ab = communities
361 .iter()
362 .find(|c| c.len() == 2)
363 .expect("应存在含 2 文件的社区");
364 let paths: Vec<String> = ab
365 .iter()
366 .map(|nid| kg.graph.node_weight(*nid).unwrap().file_path.clone().unwrap())
367 .collect();
368 assert!(paths.contains(&"src/a.rs".to_string()));
369 assert!(paths.contains(&"src/b.rs".to_string()));
370 }
371
372 #[test]
373 fn test_detect_communities_empty_graph() {
374 let kg = KnowledgeGraph::default();
375 assert!(detect_communities(&kg).is_empty());
376 }
377
378 #[test]
379 fn test_detect_communities_single_file() {
380 let mut kg = KnowledgeGraph::default();
381 let g = &mut kg.graph;
382 g.add_node(CodeNode {
383 id: NodeId::new(0),
384 kind: NodeKind::File,
385 name: "src/main.rs".into(),
386 file_path: Some("src/main.rs".into()),
387 line_range: None,
388 doc_comment: None,
389 signature: None, visibility: None,
390 module_path: vec!["src".into()],
391 });
392 let communities = detect_communities(&kg);
393 assert_eq!(communities.len(), 1);
394 assert_eq!(communities[0].len(), 1);
395 }
396
397 #[test]
398 fn test_community_name_common_prefix() {
399 let files = vec!["src/net/tcp.rs".to_string(), "src/net/udp.rs".to_string()];
400 assert_eq!(community_name(&files, 0), "src::net");
401 }
402
403 #[test]
404 fn test_community_name_single_file() {
405 let files = vec!["src/config.rs".to_string()];
406 assert_eq!(community_name(&files, 3), "src");
407 }
408
409 #[test]
410 fn test_community_name_most_populated_dir() {
411 let files = vec![
413 "app/main.rs".to_string(),
414 "app/util.rs".to_string(),
415 "lib/helper.rs".to_string(),
416 ];
417 assert_eq!(community_name(&files, 1), "app");
418 }
419
420 #[test]
421 fn test_community_name_fallback() {
422 let files = vec!["main.rs".to_string()];
423 assert_eq!(community_name(&files, 7), "module_7");
424 assert_eq!(community_name(&[], 2), "module_2");
425 }
426
427 #[test]
430 fn test_detect_communities_stable_order() {
431 let mut kg = KnowledgeGraph::default();
432 let g = &mut kg.graph;
433 let add_comm = |g: &mut petgraph::stable_graph::StableDiGraph<CodeNode, CodeEdge>,
435 path: &str| {
436 let nid = g.add_node(CodeNode {
437 id: NodeId::new(g.node_count()),
438 kind: NodeKind::File,
439 name: path.into(),
440 file_path: Some(path.into()),
441 line_range: None,
442 doc_comment: None,
443 signature: None, visibility: None,
444 module_path: dir_segments(path),
445 });
446 let eid = g.add_node(CodeNode {
447 id: NodeId::new(g.node_count()),
448 kind: NodeKind::Function,
449 name: format!("f{}", nid.index()),
450 file_path: Some(path.into()),
451 line_range: None,
452 doc_comment: None,
453 signature: None, visibility: None,
454 module_path: Vec::new(),
455 });
456 g.add_edge(
457 nid,
458 eid,
459 CodeEdge {
460 id: EdgeId::new(g.edge_count()),
461 kind: EdgeKind::Contains,
462 source: nid,
463 target: eid,
464 weight: 1.0,
465 location: None,
466 },
467 );
468 };
469 add_comm(g, "src/net/tcp.rs");
470 add_comm(g, "src/net/udp.rs");
471 add_comm(g, "src/util.rs");
473 let f = |g: &mut petgraph::stable_graph::StableDiGraph<CodeNode, CodeEdge>, path: &str| {
475 g.node_indices()
476 .find(|n| g.node_weight(*n).map(|n| n.name == path).unwrap_or(false))
477 .unwrap()
478 };
479 let _tcp = f(g, "src/net/tcp.rs");
480 let _udp = f(g, "src/net/udp.rs");
481 let fns: Vec<_> = g
483 .node_indices()
484 .filter(|n| g.node_weight(*n).map(|n| n.kind == NodeKind::Function).unwrap_or(false))
485 .collect();
486 g.add_edge(
487 fns[0],
488 fns[1],
489 CodeEdge {
490 id: EdgeId::new(g.edge_count()),
491 kind: EdgeKind::Calls,
492 source: fns[0],
493 target: fns[1],
494 weight: 0.7,
495 location: None,
496 },
497 );
498 let communities = detect_communities(&kg);
499 assert_eq!(communities.len(), 2);
500 assert_eq!(communities[0].len(), 2, "大社区应排前(稳定重排序)");
502 assert_eq!(communities[1].len(), 1);
503 }
504
505 fn make_dirs_graph(
510 n_dirs: usize,
511 files_per_dir: usize,
512 connected_pairs: &[(usize, usize)],
513 ) -> KnowledgeGraph {
514 let mut kg = KnowledgeGraph::default();
515 let g = &mut kg.graph;
516 let mut first_entity: HashMap<String, NodeId> = HashMap::new();
518 for d in 0..n_dirs {
519 let dir = format!("dir{d:02}");
520 for f in 0..files_per_dir {
521 let path = format!("{dir}/f{f}.rs");
522 let nid = g.add_node(CodeNode {
523 id: NodeId::new(g.node_count()),
524 kind: NodeKind::File,
525 name: path.clone(),
526 file_path: Some(path.clone()),
527 line_range: None,
528 doc_comment: None,
529 signature: None,
530 visibility: None,
531 module_path: dir_segments(&path),
532 });
533 let eid = g.add_node(CodeNode {
534 id: NodeId::new(g.node_count()),
535 kind: NodeKind::Function,
536 name: format!("f{f}"),
537 file_path: Some(path),
538 line_range: None,
539 doc_comment: None,
540 signature: None,
541 visibility: None,
542 module_path: Vec::new(),
543 });
544 g.add_edge(
545 nid,
546 eid,
547 CodeEdge {
548 id: EdgeId::new(g.edge_count()),
549 kind: EdgeKind::Contains,
550 source: nid,
551 target: eid,
552 weight: 1.0,
553 location: None,
554 },
555 );
556 first_entity.entry(dir.clone()).or_insert(eid);
557 }
558 }
559 for (a, b) in connected_pairs {
560 let ea = first_entity[&format!("dir{a:02}")];
561 let eb = first_entity[&format!("dir{b:02}")];
562 g.add_edge(
563 ea,
564 eb,
565 CodeEdge {
566 id: EdgeId::new(g.edge_count()),
567 kind: EdgeKind::Calls,
568 source: ea,
569 target: eb,
570 weight: 0.7,
571 location: None,
572 },
573 );
574 }
575 kg
576 }
577
578 #[test]
583 fn test_detect_communities_entity_level_below_threshold() {
584 let kg = make_dirs_graph(20, 2, &[(0, 1), (1, 2), (3, 4)]);
586
587 let first = detect_communities(&kg);
588 let second = detect_communities(&kg);
589 assert_eq!(first, second, "同图两次划分必须完全一致(确定性)");
590
591 assert!(
592 first.len() > 20,
593 "实体级划分社区数应多于目录数(独立目录每文件一社区), 实际 {} 个社区",
594 first.len()
595 );
596 let mixed = first.iter().any(|c| {
598 let dirs_in: std::collections::HashSet<&str> = c
599 .iter()
600 .filter_map(|nid| kg.graph.node_weight(*nid).and_then(|n| n.file_path.as_deref()))
601 .filter_map(|p| p.rsplit_once('/').map(|(d, _)| d))
602 .collect();
603 dirs_in.len() > 1
604 });
605 assert!(mixed, "跨目录调用链应产生混合目录社区, 实际: {:?}", first);
606 }
607
608 #[test]
613 fn test_detect_communities_dir_level_at_and_above_threshold() {
614 for (n_dirs, files_per_dir) in [(24usize, 3usize), (30, 2), (40, 3)] {
616 let pairs: Vec<(usize, usize)> = (0..n_dirs - 1).map(|a| (a, a + 1)).collect();
618 let kg = make_dirs_graph(n_dirs, files_per_dir, &pairs);
619
620 let first = detect_communities(&kg);
621 let second = detect_communities(&kg);
622 assert_eq!(first, second, "{n_dirs} 目录两次划分必须一致(确定性)");
623 assert_eq!(
624 first.len(),
625 n_dirs,
626 "{n_dirs} 目录应产出 {n_dirs} 个社区, 实际 {}",
627 first.len()
628 );
629
630 for comm in &first {
631 let mut paths: Vec<String> = comm
632 .iter()
633 .filter_map(|nid| kg.graph.node_weight(*nid).and_then(|n| n.file_path.clone()))
634 .collect();
635 paths.sort();
636 assert_eq!(paths.len(), files_per_dir, "每社区应为单目录全部文件: {paths:?}");
637 let dirs_in: std::collections::HashSet<&str> = paths
638 .iter()
639 .filter_map(|p| p.rsplit_once('/').map(|(d, _)| d))
640 .collect();
641 assert_eq!(dirs_in.len(), 1, "社区内文件必须同属一个目录: {paths:?}");
642 }
643 }
644 }
645
646 #[test]
652 fn test_detect_communities_single_dir_repo() {
653 let kg = make_dirs_graph(1, 10, &[]);
655 let first = detect_communities(&kg);
656 let second = detect_communities(&kg);
657 assert_eq!(first, second, "单目录仓库两次划分必须一致(确定性)");
658 assert_eq!(
659 first.len(),
660 1,
661 "单目录仓库应产出 1 个社区(整库一个模块), 实际 {} 个",
662 first.len()
663 );
664 assert_eq!(first[0].len(), 10, "社区应包含全部 10 个文件");
665 let all_in_dir00 = first[0].iter().all(|nid| {
666 kg.graph
667 .node_weight(*nid)
668 .and_then(|n| n.file_path.as_deref())
669 .is_some_and(|p| p.starts_with("dir00/"))
670 });
671 assert!(all_in_dir00, "社区内所有文件必须同属唯一目录");
672
673 let mut kg2 = KnowledgeGraph::default();
675 let g = &mut kg2.graph;
676 for i in 0..5 {
677 let path = format!("main{i}.rs");
678 let nid = g.add_node(CodeNode {
679 id: NodeId::new(g.node_count()),
680 kind: NodeKind::File,
681 name: path.clone(),
682 file_path: Some(path),
683 line_range: None,
684 doc_comment: None,
685 signature: None,
686 visibility: None,
687 module_path: Vec::new(),
688 });
689 let eid = g.add_node(CodeNode {
690 id: NodeId::new(g.node_count()),
691 kind: NodeKind::Function,
692 name: format!("f{i}"),
693 file_path: None,
694 line_range: None,
695 doc_comment: None,
696 signature: None,
697 visibility: None,
698 module_path: Vec::new(),
699 });
700 g.add_edge(
701 nid,
702 eid,
703 CodeEdge {
704 id: EdgeId::new(g.edge_count()),
705 kind: EdgeKind::Contains,
706 source: nid,
707 target: eid,
708 weight: 1.0,
709 location: None,
710 },
711 );
712 }
713 let comms = detect_communities(&kg2);
714 assert_eq!(comms.len(), 1, "根目录散文件仓库也应整体一社区");
715 assert_eq!(comms[0].len(), 5);
716 }
717}