1use crate::error::{ForgeError, Result};
2use crate::knowledge::types::{self, CfgBlockData, GraphNode};
3use crate::knowledge::KnowledgeGraph;
4
5#[derive(Clone, Debug)]
8pub struct SourceSpan {
9 pub file: String,
10 pub line: usize,
11 pub byte_start: u32,
12 pub byte_end: u32,
13}
14
15impl SourceSpan {
16 pub fn new(file: impl Into<String>, line: usize, byte_start: u32, byte_end: u32) -> Self {
17 Self {
18 file: file.into(),
19 line,
20 byte_start,
21 byte_end,
22 }
23 }
24}
25
26impl KnowledgeGraph {
27 pub fn get_node(&self, node_id: i64) -> Result<GraphNode> {
28 let entity = self
29 .backend
30 .get_node(Self::snapshot(), node_id)
31 .map_err(|e| ForgeError::DatabaseError(format!("Node not found: {}", e)))?;
32 Ok(GraphNode {
33 id: node_id,
34 kind: entity.kind,
35 name: entity.name,
36 file_path: entity.file_path,
37 data: entity.data,
38 })
39 }
40
41 pub fn find_nodes_by_kind(&self, kind: &str) -> Result<Vec<GraphNode>> {
42 let snap = Self::snapshot();
43 let ids = self
44 .backend
45 .entity_ids()
46 .map_err(|e| ForgeError::DatabaseError(format!("Entity list failed: {}", e)))?;
47 let mut results = Vec::new();
48 for id in ids {
49 if let Ok(entity) = self.backend.get_node(snap, id) {
50 if entity.kind == kind {
51 results.push(GraphNode {
52 id,
53 kind: entity.kind,
54 name: entity.name,
55 file_path: entity.file_path,
56 data: entity.data,
57 });
58 }
59 }
60 }
61 Ok(results)
62 }
63
64 pub fn add_symbol(
65 &self,
66 name: &str,
67 symbol_kind: &str,
68 qualified_name: &str,
69 span: &SourceSpan,
70 language: &str,
71 parent_id: Option<i64>,
72 ) -> Result<i64> {
73 let mut data = serde_json::json!({
74 "symbol_kind": symbol_kind,
75 "qualified_name": qualified_name,
76 "file": span.file,
77 "line": span.line,
78 "byte_start": span.byte_start,
79 "byte_end": span.byte_end,
80 "language": language,
81 });
82 if let Some(pid) = parent_id {
83 data["parent_id"] = serde_json::json!(pid);
84 }
85 self.insert_node(types::node::SYMBOL, name, Some(&span.file), data)
86 }
87
88 pub fn add_file(&self, path: &str, language: &str, hash: &str) -> Result<i64> {
89 let data = serde_json::json!({
90 "path": path,
91 "language": language,
92 "hash": hash,
93 });
94 self.insert_node(types::node::FILE, path, None, data)
95 }
96
97 pub fn add_discovery(
98 &self,
99 agent: &str,
100 discovery_type: &str,
101 target: &str,
102 metadata: serde_json::Value,
103 ) -> Result<i64> {
104 let data = serde_json::json!({
105 "discovery_type": discovery_type,
106 "agent": agent,
107 "timestamp": chrono::Utc::now().to_rfc3339(),
108 "metadata": metadata,
109 });
110 self.insert_node(types::node::DISCOVERY, target, None, data)
111 }
112
113 pub fn add_issue(
114 &self,
115 severity: &str,
116 description: &str,
117 rule_id: Option<&str>,
118 ) -> Result<i64> {
119 let mut data = serde_json::json!({"severity": severity, "description": description,});
120 if let Some(rid) = rule_id {
121 data["rule_id"] = serde_json::json!(rid);
122 }
123 self.insert_node(types::node::ISSUE, description, None, data)
124 }
125
126 pub fn add_pattern(
127 &self,
128 pattern_type: &str,
129 confidence: f64,
130 description: &str,
131 ) -> Result<i64> {
132 let data = serde_json::json!({
133 "pattern_type": pattern_type,
134 "confidence": confidence,
135 "description": description,
136 });
137 self.insert_node(types::node::PATTERN, pattern_type, None, data)
138 }
139
140 pub fn add_knowledge(
141 &self,
142 source: &str,
143 title: &str,
144 tags: &[String],
145 summary: &str,
146 ) -> Result<i64> {
147 let data = serde_json::json!({
148 "source": source,
149 "title": title,
150 "tags": tags,
151 "summary": summary,
152 });
153 self.insert_node(types::node::KNOWLEDGE, title, None, data)
154 }
155
156 pub fn add_hotspot(
157 &self,
158 complexity: u32,
159 risk_score: f64,
160 loop_depth: u32,
161 description: &str,
162 ) -> Result<i64> {
163 let data = serde_json::json!({
164 "complexity": complexity,
165 "risk_score": risk_score,
166 "loop_depth": loop_depth,
167 "description": description,
168 });
169 self.insert_node(types::node::HOTSPOT, description, None, data)
170 }
171
172 pub fn add_cfg_block(&self, function_id: i64, block: &CfgBlockData) -> Result<i64> {
173 let data = serde_json::json!({
174 "function_id": function_id,
175 "start_byte": block.start_byte,
176 "end_byte": block.end_byte,
177 "block_kind": block.block_kind,
178 "is_error": block.is_error,
179 });
180 self.insert_node(
181 types::node::CFG_BLOCK,
182 &format!("block_{}", block.start_byte),
183 None,
184 data,
185 )
186 }
187}
188
189#[cfg(test)]
190mod tests {
191 use crate::knowledge::{open_kg, CfgBlockData, SourceSpan};
192
193 #[test]
194 fn test_add_symbol_node() {
195 let (_temp, kg) = open_kg();
196 let id = kg
197 .add_symbol(
198 "my_func",
199 "Function",
200 "crate::module::my_func",
201 &SourceSpan::new("src/lib.rs", 42, 100, 200),
202 "Rust",
203 None,
204 )
205 .expect("invariant: fresh graph accepts inserts");
206 assert!(id > 0);
207 let node = kg
208 .get_node(id)
209 .expect("invariant: just-inserted node is retrievable");
210 assert_eq!(node.kind, "symbol");
211 assert_eq!(node.name, "my_func");
212 assert_eq!(node.prop_str("symbol_kind"), Some("Function"));
213 assert_eq!(
214 node.prop_str("qualified_name"),
215 Some("crate::module::my_func")
216 );
217 assert_eq!(node.prop_str("file"), Some("src/lib.rs"));
218 assert_eq!(node.prop_u64("line"), Some(42));
219 }
220
221 #[test]
222 fn test_add_file_node() {
223 let (_temp, kg) = open_kg();
224 let id = kg
225 .add_file("src/lib.rs", "Rust", "abc123")
226 .expect("invariant: fresh graph accepts inserts");
227 let node = kg
228 .get_node(id)
229 .expect("invariant: just-inserted node is retrievable");
230 assert_eq!(node.kind, "file");
231 assert_eq!(node.name, "src/lib.rs");
232 assert_eq!(node.prop_str("language"), Some("Rust"));
233 assert_eq!(node.prop_str("hash"), Some("abc123"));
234 }
235
236 #[test]
237 fn test_add_discovery_node() {
238 let (_temp, kg) = open_kg();
239 let id = kg
240 .add_discovery(
241 "claude1",
242 "Symbol",
243 "my_func",
244 serde_json::json!({"complexity": 8}),
245 )
246 .expect("invariant: fresh graph accepts inserts");
247 let node = kg
248 .get_node(id)
249 .expect("invariant: just-inserted node is retrievable");
250 assert_eq!(node.kind, "discovery");
251 assert_eq!(node.name, "my_func");
252 assert_eq!(node.prop_str("agent"), Some("claude1"));
253 assert_eq!(node.prop_str("discovery_type"), Some("Symbol"));
254 }
255
256 #[test]
257 fn test_add_issue_node() {
258 let (_temp, kg) = open_kg();
259 let id = kg
260 .add_issue("high", "unwrap in production code", Some("M001"))
261 .expect("invariant: fresh graph accepts inserts");
262 let node = kg
263 .get_node(id)
264 .expect("invariant: just-inserted node is retrievable");
265 assert_eq!(node.kind, "issue");
266 assert_eq!(node.prop_str("severity"), Some("high"));
267 assert_eq!(node.prop_str("rule_id"), Some("M001"));
268 }
269
270 #[test]
271 fn test_add_pattern_node() {
272 let (_temp, kg) = open_kg();
273 let id = kg
274 .add_pattern("builder", 0.92, "builder pattern detected")
275 .expect("invariant: fresh graph accepts inserts");
276 let node = kg
277 .get_node(id)
278 .expect("invariant: just-inserted node is retrievable");
279 assert_eq!(node.kind, "pattern");
280 assert_eq!(node.prop_str("pattern_type"), Some("builder"));
281 assert_eq!(node.prop_f64("confidence"), Some(0.92));
282 }
283
284 #[test]
285 fn test_add_knowledge_node() {
286 let (_temp, kg) = open_kg();
287 let tags = vec!["auth".to_string(), "middleware".to_string()];
288 let id = kg
289 .add_knowledge("wiki", "Auth Architecture", &tags, "Overview of auth flow")
290 .expect("invariant: fresh graph accepts inserts");
291 let node = kg
292 .get_node(id)
293 .expect("invariant: just-inserted node is retrievable");
294 assert_eq!(node.kind, "knowledge");
295 assert_eq!(node.prop_str("source"), Some("wiki"));
296 assert_eq!(node.prop_str("title"), Some("Auth Architecture"));
297 }
298
299 #[test]
300 fn test_add_hotspot_node() {
301 let (_temp, kg) = open_kg();
302 let id = kg
303 .add_hotspot(15, 0.85, 3, "high complexity loop")
304 .expect("invariant: fresh graph accepts inserts");
305 let node = kg
306 .get_node(id)
307 .expect("invariant: just-inserted node is retrievable");
308 assert_eq!(node.kind, "hotspot");
309 assert_eq!(node.prop_u64("complexity"), Some(15));
310 assert_eq!(node.prop_f64("risk_score"), Some(0.85));
311 }
312
313 #[test]
314 fn test_add_cfg_block_node() {
315 let (_temp, kg) = open_kg();
316 let block = CfgBlockData {
317 start_byte: 100,
318 end_byte: 200,
319 block_kind: "Basic".to_string(),
320 is_error: false,
321 };
322 let id = kg
323 .add_cfg_block(42, &block)
324 .expect("invariant: fresh graph accepts inserts");
325 let node = kg
326 .get_node(id)
327 .expect("invariant: just-inserted node is retrievable");
328 assert_eq!(node.kind, "cfg_block");
329 assert_eq!(node.prop_u64("function_id"), Some(42));
330 assert_eq!(node.prop_str("block_kind"), Some("Basic"));
331 }
332
333 #[test]
334 fn test_find_nodes_by_kind() {
335 let (_temp, kg) = open_kg();
336 kg.add_symbol(
337 "func_a",
338 "Function",
339 "a",
340 &SourceSpan::new("f.rs", 1, 0, 10),
341 "Rust",
342 None,
343 )
344 .expect("invariant: fresh graph accepts inserts");
345 kg.add_symbol(
346 "func_b",
347 "Function",
348 "b",
349 &SourceSpan::new("f.rs", 2, 0, 10),
350 "Rust",
351 None,
352 )
353 .expect("invariant: fresh graph accepts inserts");
354 kg.add_file("f.rs", "Rust", "hash")
355 .expect("invariant: fresh graph accepts inserts");
356
357 let symbols = kg
358 .find_nodes_by_kind("symbol")
359 .expect("invariant: query on valid graph succeeds");
360 assert_eq!(symbols.len(), 2);
361 let files = kg
362 .find_nodes_by_kind("file")
363 .expect("invariant: query on valid graph succeeds");
364 assert_eq!(files.len(), 1);
365 }
366
367 #[test]
368 fn test_get_node_not_found() {
369 let (_temp, kg) = open_kg();
370 let result = kg.get_node(99999);
371 assert!(result.is_err());
372 }
373}