Skip to main content

graphify_core/
model.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::confidence::Confidence;
6
7// ---------------------------------------------------------------------------
8// NodeType
9// ---------------------------------------------------------------------------
10
11/// The kind of entity a graph node represents.
12///
13/// Serialized as lowercase strings (e.g. `"class"`, `"function"`).
14#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
15#[serde(rename_all = "lowercase")]
16pub enum NodeType {
17    Class,
18    Function,
19    Module,
20    Concept,
21    Paper,
22    Image,
23    File,
24    Method,
25    Interface,
26    Enum,
27    Struct,
28    Trait,
29    Constant,
30    Variable,
31    Package,
32    Namespace,
33}
34
35// ---------------------------------------------------------------------------
36// GraphNode
37// ---------------------------------------------------------------------------
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct GraphNode {
41    pub id: String,
42    pub label: String,
43    pub source_file: String,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub source_location: Option<String>,
46    pub node_type: NodeType,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub community: Option<usize>,
49    #[serde(flatten)]
50    pub extra: HashMap<String, serde_json::Value>,
51}
52
53// ---------------------------------------------------------------------------
54// GraphEdge
55// ---------------------------------------------------------------------------
56
57fn default_confidence_score() -> f64 {
58    1.0
59}
60
61fn default_weight() -> f64 {
62    1.0
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct GraphEdge {
67    pub source: String,
68    pub target: String,
69    pub relation: String,
70    pub confidence: Confidence,
71    #[serde(default = "default_confidence_score")]
72    pub confidence_score: f64,
73    pub source_file: String,
74    #[serde(skip_serializing_if = "Option::is_none")]
75    pub source_location: Option<String>,
76    #[serde(default = "default_weight")]
77    pub weight: f64,
78    #[serde(flatten)]
79    pub extra: HashMap<String, serde_json::Value>,
80}
81
82// ---------------------------------------------------------------------------
83// Hyperedge
84// ---------------------------------------------------------------------------
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct Hyperedge {
88    pub nodes: Vec<String>,
89    pub relation: String,
90    pub label: String,
91}
92
93// ---------------------------------------------------------------------------
94// ExtractionResult
95// ---------------------------------------------------------------------------
96
97#[derive(Debug, Clone, Default, Serialize, Deserialize)]
98pub struct ExtractionResult {
99    pub nodes: Vec<GraphNode>,
100    pub edges: Vec<GraphEdge>,
101    #[serde(default)]
102    pub hyperedges: Vec<Hyperedge>,
103}
104
105// ---------------------------------------------------------------------------
106// CommunityInfo
107// ---------------------------------------------------------------------------
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct CommunityInfo {
111    pub id: usize,
112    pub nodes: Vec<String>,
113    pub cohesion: f64,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub label: Option<String>,
116}
117
118// ---------------------------------------------------------------------------
119// AnalysisResult & helpers
120// ---------------------------------------------------------------------------
121
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct GodNode {
124    pub id: String,
125    pub label: String,
126    pub degree: usize,
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub community: Option<usize>,
129}
130
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct Surprise {
133    pub source: String,
134    pub target: String,
135    pub source_community: usize,
136    pub target_community: usize,
137    pub relation: String,
138}
139
140#[derive(Debug, Clone, Default, Serialize, Deserialize)]
141pub struct AnalysisResult {
142    pub god_nodes: Vec<GodNode>,
143    pub surprises: Vec<Surprise>,
144    pub questions: Vec<String>,
145}
146
147// ---------------------------------------------------------------------------
148// Tests
149// ---------------------------------------------------------------------------
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    fn sample_node() -> GraphNode {
156        GraphNode {
157            id: "my_class".into(),
158            label: "MyClass".into(),
159            source_file: "src/main.rs".into(),
160            source_location: Some("L42".into()),
161            node_type: NodeType::Class,
162            community: None,
163            extra: HashMap::new(),
164        }
165    }
166
167    fn sample_edge() -> GraphEdge {
168        GraphEdge {
169            source: "a".into(),
170            target: "b".into(),
171            relation: "calls".into(),
172            confidence: Confidence::Extracted,
173            confidence_score: 1.0,
174            source_file: "src/main.rs".into(),
175            source_location: None,
176            weight: 1.0,
177            extra: HashMap::new(),
178        }
179    }
180
181    #[test]
182    fn node_type_serializes_lowercase() {
183        assert_eq!(
184            serde_json::to_string(&NodeType::Class).unwrap(),
185            r#""class""#
186        );
187        assert_eq!(
188            serde_json::to_string(&NodeType::Function).unwrap(),
189            r#""function""#
190        );
191        assert_eq!(
192            serde_json::to_string(&NodeType::Namespace).unwrap(),
193            r#""namespace""#
194        );
195    }
196
197    #[test]
198    fn node_roundtrip() {
199        let node = sample_node();
200        let json = serde_json::to_string(&node).unwrap();
201        let back: GraphNode = serde_json::from_str(&json).unwrap();
202        assert_eq!(back.id, "my_class");
203        assert_eq!(back.node_type, NodeType::Class);
204    }
205
206    #[test]
207    fn node_skip_none_fields() {
208        let mut node = sample_node();
209        node.source_location = None;
210        node.community = None;
211        let json = serde_json::to_string(&node).unwrap();
212        assert!(!json.contains("source_location"));
213        assert!(!json.contains("community"));
214    }
215
216    #[test]
217    fn edge_defaults() {
218        let json = r#"{
219            "source": "a",
220            "target": "b",
221            "relation": "calls",
222            "confidence": "EXTRACTED",
223            "source_file": "x.rs"
224        }"#;
225        let edge: GraphEdge = serde_json::from_str(json).unwrap();
226        assert!((edge.confidence_score - 1.0).abs() < f64::EPSILON);
227        assert!((edge.weight - 1.0).abs() < f64::EPSILON);
228    }
229
230    #[test]
231    fn edge_roundtrip() {
232        let edge = sample_edge();
233        let json = serde_json::to_string(&edge).unwrap();
234        let back: GraphEdge = serde_json::from_str(&json).unwrap();
235        assert_eq!(back.relation, "calls");
236        assert_eq!(back.confidence, Confidence::Extracted);
237    }
238
239    #[test]
240    fn extraction_result_default() {
241        let r = ExtractionResult::default();
242        assert!(r.nodes.is_empty());
243        assert!(r.edges.is_empty());
244        assert!(r.hyperedges.is_empty());
245    }
246
247    #[test]
248    fn extra_fields_flatten() {
249        let mut node = sample_node();
250        node.extra
251            .insert("custom".into(), serde_json::Value::Bool(true));
252        let json = serde_json::to_string(&node).unwrap();
253        assert!(json.contains(r#""custom":true"#));
254    }
255
256    #[test]
257    fn community_info_roundtrip() {
258        let ci = CommunityInfo {
259            id: 0,
260            nodes: vec!["a".into(), "b".into()],
261            cohesion: 0.85,
262            label: Some("cluster-0".into()),
263        };
264        let json = serde_json::to_string(&ci).unwrap();
265        let back: CommunityInfo = serde_json::from_str(&json).unwrap();
266        assert_eq!(back.id, 0);
267        assert_eq!(back.nodes.len(), 2);
268    }
269
270    #[test]
271    fn analysis_result_default() {
272        let ar = AnalysisResult::default();
273        assert!(ar.god_nodes.is_empty());
274        assert!(ar.surprises.is_empty());
275        assert!(ar.questions.is_empty());
276    }
277}