code_repo_wiki/generate/
chunk.rs

1use std::collections::{HashMap, HashSet};
2use std::path::PathBuf;
3
4use petgraph::stable_graph::NodeIndex;
5use petgraph::visit::EdgeRef;
6
7use crate::ingest::parser::{Entity, FileInsight, ImportStmt};
8use crate::model::{EdgeKind, KnowledgeGraph, ModuleCluster, NodeKind};
9
10/// AST 感知的数据块，用于 LLM 生成
11///
12/// 每个 Chunk 代表一个生成单元——可以是一个模块或单个文件。
13#[derive(Debug, Clone)]
14pub struct Chunk {
15    /// 模块路径（如 ["crate", "generate", "llm"]）
16    pub module_path: Vec<String>,
17    /// 块中的实体列表
18    pub entities: Vec<Entity>,
19    /// 块的导入语句
20    pub imports: Vec<ImportStmt>,
21    /// 依赖的其他模块名
22    pub dependencies: Vec<String>,
23    /// 关联的源文件路径
24    pub file_paths: Vec<PathBuf>,
25    /// 每个实体所属的源文件（与 entities 平行；空表示未记录，
26    /// 供增量场景按实体级过滤摘要生成——演进计划 T2.3）
27    pub entity_sources: Vec<PathBuf>,
28}
29
30impl Chunk {
31    /// 是否为空块（无实体和导入）
32    pub fn is_empty(&self) -> bool {
33        self.entities.is_empty() && self.imports.is_empty()
34    }
35
36    /// 实体数量
37    pub fn entity_count(&self) -> usize {
38        self.entities.len()
39    }
40}
41
42/// 从知识图谱构建 NodeId → 文件路径的映射
43///
44/// 只包含 File 类型的节点，用于 chunk_by_module 的精确路径匹配。
45pub fn build_node_to_file_map(graph: &KnowledgeGraph) -> HashMap<NodeIndex, PathBuf> {
46    let mut map = HashMap::new();
47    for n in graph.graph.node_indices() {
48        if let Some(w) = graph.graph.node_weight(n)
49            && w.kind == NodeKind::File
50            && let Some(ref fp) = w.file_path
51        {
52            map.insert(n, PathBuf::from(fp));
53        }
54    }
55    map
56}
57
58/// 以模块为单位进行 AST 感知分块
59///
60/// 将 FileInsight 按 ModuleCluster 分组，通过 node_id → file_path 映射精确匹配，
61/// 同一模块的多个文件合并到一个 Chunk。
62/// 返回的 Chunk 列表与 modules 一一对应。
63pub fn chunk_by_module(
64    insights: &[FileInsight],
65    modules: &[ModuleCluster],
66    graph: &KnowledgeGraph,
67) -> Vec<Chunk> {
68    let node_to_file = build_node_to_file_map(graph);
69    let mut chunks = Vec::with_capacity(modules.len());
70
71    // 预计算模块名 → 节点集合的映射，用于依赖分析
72    let module_node_ids: std::collections::HashMap<&str, HashSet<NodeIndex>> = modules
73        .iter()
74        .map(|m| (m.name.as_str(), m.node_ids.iter().copied().collect()))
75        .collect();
76
77    for module in modules {
78        // 收集该模块中所有 File 节点的路径
79        let module_file_paths: HashSet<&PathBuf> = module
80            .node_ids
81            .iter()
82            .filter_map(|nid| node_to_file.get(nid))
83            .collect();
84
85        let mut entities = Vec::new();
86        let mut imports = Vec::new();
87        let mut file_paths = Vec::new();
88        let mut entity_sources = Vec::new();
89
90        for insight in insights {
91            if module_file_paths.contains(&insight.path) {
92                // 记录实体 → 源文件归属（与 entities 平行，供 T2.3 实体级过滤）
93                for _ in &insight.entities {
94                    entity_sources.push(insight.path.clone());
95                }
96                entities.extend(insight.entities.clone());
97                imports.extend(insight.imports.clone());
98                if !file_paths.contains(&insight.path) {
99                    file_paths.push(insight.path.clone());
100                }
101            }
102        }
103
104        // 实体与文件归属按同一键（name）排序去重，保证两者仍平行
105        let mut paired: Vec<(Entity, PathBuf)> = entities
106            .into_iter()
107            .zip(entity_sources)
108            .collect();
109        paired.sort_by(|a, b| a.0.name.cmp(&b.0.name));
110        // N4：同名实体去重——不同文件定义同名实体时按排序后首个保留，
111        // 此前静默丢弃无任何提示；告警暴露去重事实（名称冲突常见于
112        // 重载/同名导出，模块页与搜索索引只保留一个定义）
113        let before = paired.len();
114        paired.dedup_by(|a, b| a.0.name == b.0.name);
115        if paired.len() < before {
116            tracing::warn!(
117                "模块 {} 去重 {} 个同名实体（保留排序后首个定义）",
118                module.name,
119                before - paired.len()
120            );
121        }
122        let entities: Vec<Entity> = paired.iter().map(|(e, _)| e.clone()).collect();
123        let entity_sources: Vec<PathBuf> = paired.iter().map(|(_, f)| f.clone()).collect();
124
125        let module_path: Vec<String> = module.name.split("::").map(|s| s.to_string()).collect();
126
127        // 计算实际依赖：从本模块节点出发，通过 Imports 边到达其他模块的节点
128        // module_node_ids 是 HashMap，迭代序随机——deps 必须排序，
129        // 否则卡片 frontmatter 的 dependencies 顺序跨次漂移（确定性评测失败）
130        let mut deps: Vec<String> = Vec::new();
131        for (&other_name, other_set) in &module_node_ids {
132            if other_name == module.name {
133                continue;
134            }
135            let has_dep = module.node_ids.iter().any(|nid| {
136                graph.graph.edges(*nid).any(|e| {
137                    let kind = &graph.graph[e.id()].kind;
138                    kind == &EdgeKind::Imports && other_set.contains(&e.target())
139                })
140            });
141            if has_dep {
142                deps.push(other_name.to_string());
143            }
144        }
145        deps.sort();
146
147        chunks.push(Chunk {
148            module_path,
149            entities,
150            imports,
151            dependencies: deps,
152            file_paths,
153            entity_sources,
154        });
155    }
156
157    chunks
158}
159
160/// 以文件为单位分块（适用于 Level 0，无模块聚类信息时回退）
161pub fn chunk_by_file(insight: &FileInsight) -> Chunk {
162    let module_path: Vec<String> = insight
163        .path
164        .parent()
165        .and_then(|p| {
166            // 只取普通目录组件，过滤盘符（Prefix）/根目录（RootDir）等
167            // 否则 Windows 绝对路径会生成含 ":\" 的 module_path，导致 wiki 文件名非法
168            p.components()
169                .filter(|c| matches!(c, std::path::Component::Normal(_)))
170                .map(|c| c.as_os_str().to_string_lossy().to_string())
171                .reduce(|a, b| format!("{}::{}", a, b))
172        })
173        .map(|s| s.split("::").map(|p| p.to_string()).collect())
174        .unwrap_or_default();
175
176    Chunk {
177        module_path,
178        entities: insight.entities.clone(),
179        imports: insight.imports.clone(),
180        dependencies: Vec::new(),
181        file_paths: vec![insight.path.clone()],
182        entity_sources: insight.entities.iter().map(|_| insight.path.clone()).collect(),
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use crate::ingest::parser::Entity;
190    use crate::model::KnowledgeGraph;
191
192    fn make_entity(name: &str, kind: &str) -> Entity {
193        Entity {
194            name: name.to_string(),
195            kind: kind.to_string(),
196            line_start: 1,
197            line_end: 10,
198            doc_comment: None,
199            signature: None, visibility: None,
200        }
201    }
202
203    fn make_insight(file_name: &str, entities: Vec<Entity>) -> FileInsight {
204        FileInsight {
205            path: PathBuf::from(file_name),
206            language: "rust".into(),
207            entities,
208            imports: Vec::new(),
209            doc_comments: Vec::new(),
210            source: String::new(),
211        }
212    }
213
214    #[test]
215    fn test_chunk_by_file() {
216        let entity = make_entity("MyStruct", "struct");
217        let insight = make_insight("src/lib.rs", vec![entity]);
218        let chunk = chunk_by_file(&insight);
219
220        assert_eq!(chunk.entity_count(), 1);
221        assert!(!chunk.module_path.is_empty());
222    }
223
224    #[test]
225    fn test_empty_chunk() {
226        let entities = vec![make_entity("Foo", "fn")];
227        let insight = make_insight("src/main.rs", entities);
228
229        let chunk = chunk_by_file(&insight);
230        assert!(!chunk.is_empty());
231
232        let empty_chunk = Chunk {
233            module_path: vec![],
234            entities: vec![],
235            imports: vec![],
236            dependencies: vec![],
237            file_paths: vec![],
238            entity_sources: vec![],
239        };
240        assert!(empty_chunk.is_empty());
241    }
242
243    #[test]
244    fn test_chunk_by_module_empty_modules() {
245        let graph = KnowledgeGraph::default();
246        let insight = make_insight("src/lib.rs", vec![make_entity("Foo", "fn")]);
247        let chunks = chunk_by_module(&[insight], &[], &graph);
248        assert!(chunks.is_empty());
249    }
250
251    #[test]
252    fn test_build_node_to_file_map() {
253        let mut g = petgraph::stable_graph::StableDiGraph::<
254            crate::model::CodeNode,
255            crate::model::CodeEdge,
256        >::new();
257        let file_id = g.add_node(crate::model::CodeNode {
258            id: petgraph::stable_graph::NodeIndex::new(0),
259            kind: crate::model::NodeKind::File,
260            name: "lib.rs".into(),
261            file_path: Some("src/lib.rs".into()),
262            line_range: None,
263            doc_comment: None,
264            signature: None, visibility: None,
265            module_path: vec!["src".into()],
266        });
267        let fn_id = g.add_node(crate::model::CodeNode {
268            id: petgraph::stable_graph::NodeIndex::new(1),
269            kind: crate::model::NodeKind::Function,
270            name: "foo".into(),
271            file_path: Some("src/lib.rs".into()),
272            line_range: None,
273            doc_comment: None,
274            signature: None, visibility: None,
275            module_path: vec!["src".into(), "lib".into()],
276        });
277        let kg = KnowledgeGraph {
278            graph: g,
279            modules: vec![],
280        features: Vec::new(),
281        };
282        let map = build_node_to_file_map(&kg);
283        assert_eq!(map.len(), 1); // 只含 File 节点
284        assert!(map.contains_key(&file_id));
285        assert!(!map.contains_key(&fn_id));
286        assert_eq!(map[&file_id], std::path::PathBuf::from("src/lib.rs"));
287    }
288
289    #[test]
290    fn test_chunk_by_module_groups_entities_by_module() {
291        // 构造两个模块的图谱：src::a（file_a.rs 含实体 e1）、src::b（file_b.rs 含实体 e2），
292        // file_a 通过 Imports 边依赖 file_b，验证实体归属、依赖分析与 entity_sources 平行性
293        let mut graph = KnowledgeGraph::default();
294
295        let file_a = graph.graph.add_node(crate::model::CodeNode {
296            id: petgraph::stable_graph::NodeIndex::new(0),
297            kind: NodeKind::File,
298            name: "file_a.rs".into(),
299            file_path: Some("src/a/file_a.rs".into()),
300            line_range: None,
301            doc_comment: None,
302            signature: None, visibility: None,
303            module_path: vec!["src".into(), "a".into()],
304        });
305        let e1 = graph.graph.add_node(crate::model::CodeNode {
306            id: petgraph::stable_graph::NodeIndex::new(1),
307            kind: NodeKind::Function,
308            name: "e1".into(),
309            file_path: Some("src/a/file_a.rs".into()),
310            line_range: Some((1, 5)),
311            doc_comment: None,
312            signature: Some("fn e1()".into()), visibility: None,
313            module_path: vec!["src".into(), "a".into()],
314        });
315        let file_b = graph.graph.add_node(crate::model::CodeNode {
316            id: petgraph::stable_graph::NodeIndex::new(2),
317            kind: NodeKind::File,
318            name: "file_b.rs".into(),
319            file_path: Some("src/b/file_b.rs".into()),
320            line_range: None,
321            doc_comment: None,
322            signature: None, visibility: None,
323            module_path: vec!["src".into(), "b".into()],
324        });
325        let e2 = graph.graph.add_node(crate::model::CodeNode {
326            id: petgraph::stable_graph::NodeIndex::new(3),
327            kind: NodeKind::Function,
328            name: "e2".into(),
329            file_path: Some("src/b/file_b.rs".into()),
330            line_range: Some((1, 5)),
331            doc_comment: None,
332            signature: Some("fn e2()".into()), visibility: None,
333            module_path: vec!["src".into(), "b".into()],
334        });
335
336        // 文件包含实体；a 依赖 b（Imports 边 file_a → file_b）
337        graph.graph.add_edge(file_a, e1, crate::model::CodeEdge {
338            id: petgraph::stable_graph::EdgeIndex::new(0),
339            kind: EdgeKind::Contains,
340            source: file_a,
341            target: e1,
342            weight: 1.0,
343            location: None,
344        });
345        graph.graph.add_edge(file_b, e2, crate::model::CodeEdge {
346            id: petgraph::stable_graph::EdgeIndex::new(1),
347            kind: EdgeKind::Contains,
348            source: file_b,
349            target: e2,
350            weight: 1.0,
351            location: None,
352        });
353        graph.graph.add_edge(file_a, file_b, crate::model::CodeEdge {
354            id: petgraph::stable_graph::EdgeIndex::new(2),
355            kind: EdgeKind::Imports,
356            source: file_a,
357            target: file_b,
358            weight: 1.0,
359            location: None,
360        });
361
362        // 模块按名字典序传入：chunk 与 modules 输入一一对应（chunk_by_module 不重排序）
363        let modules = vec![
364            ModuleCluster { name: "src::a".into(), node_ids: vec![file_a, e1], cohesion: 0.9, coupling: 0.1, description: None },
365            ModuleCluster { name: "src::b".into(), node_ids: vec![file_b, e2], cohesion: 0.9, coupling: 0.1, description: None },
366        ];
367        let insights = vec![
368            make_insight("src/a/file_a.rs", vec![make_entity("e1", "fn")]),
369            make_insight("src/b/file_b.rs", vec![make_entity("e2", "fn")]),
370        ];
371
372        let chunks = chunk_by_module(&insights, &modules, &graph);
373
374        // 两个模块各生成一个 chunk，模块路径按名字拆分为 ["src", "a"] / ["src", "b"]
375        assert_eq!(chunks.len(), 2);
376        assert_eq!(chunks[0].module_path, vec!["src".to_string(), "a".to_string()]);
377        assert_eq!(chunks[1].module_path, vec!["src".to_string(), "b".to_string()]);
378        // 每个 chunk 只包含自己模块的实体
379        assert_eq!(chunks[0].entities.len(), 1);
380        assert_eq!(chunks[0].entities[0].name, "e1");
381        assert_eq!(chunks[1].entities[0].name, "e2");
382        // a 依赖 b：Imports 边被依赖分析捕获
383        assert_eq!(chunks[0].dependencies, vec!["src::b".to_string()]);
384        assert!(chunks[1].dependencies.is_empty());
385        // entity_sources 与 entities 平行对应（每个实体记录其源文件）
386        assert_eq!(chunks[0].entity_sources, vec![std::path::PathBuf::from("src/a/file_a.rs")]);
387        assert_eq!(chunks[1].entity_sources, vec![std::path::PathBuf::from("src/b/file_b.rs")]);
388    }
389}
code_repo_wiki/generate/chunk.rs

code_repo_wiki/generate/
chunk.rs
