Skip to main content

project_map_cli_rust/core/
graph.rs

1use petgraph::graph::{DiGraph, NodeIndex};
2use serde::{Serialize, Deserialize};
3use std::fs;
4use std::path::Path;
5use crate::error::Result;
6
7#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
8pub enum NodeType {
9    File,
10    Symbol,
11}
12
13#[derive(Serialize, Deserialize, Debug, Clone)]
14pub struct NodeData {
15    pub path: String,
16    pub name: String,
17    pub kind: String,
18    pub line: usize,
19    pub start_byte: usize,
20    pub end_byte: usize,
21    pub node_type: NodeType,
22}
23
24#[derive(Serialize, Deserialize, Debug, Clone)]
25pub enum EdgeType {
26    Contains,
27    Imports,
28    Calls,
29}
30
31#[derive(Serialize, Deserialize)]
32struct SerializableGraph {
33    nodes: Vec<NodeData>,
34    edges: Vec<(usize, usize, EdgeType)>,
35}
36
37pub struct ProjectGraph {
38    pub graph: DiGraph<NodeData, EdgeType>,
39}
40
41impl ProjectGraph {
42    pub fn new() -> Self {
43        Self {
44            graph: DiGraph::new(),
45        }
46    }
47
48    pub fn add_node(&mut self, data: NodeData) -> NodeIndex {
49        self.graph.add_node(data)
50    }
51
52    pub fn add_edge(&mut self, from: NodeIndex, to: NodeIndex, edge_type: EdgeType) {
53        self.graph.add_edge(from, to, edge_type);
54    }
55
56    pub fn save(&self, path: &Path) -> Result<()> {
57        let mut nodes = Vec::new();
58        for i in 0..self.graph.node_count() {
59            nodes.push(self.graph[NodeIndex::new(i)].clone());
60        }
61
62        let mut edges = Vec::new();
63        for edge in self.graph.edge_indices() {
64            let (from, to) = self.graph.edge_endpoints(edge).unwrap();
65            edges.push((from.index(), to.index(), self.graph[edge].clone()));
66        }
67
68        let serializable = SerializableGraph { nodes, edges };
69        let json = serde_json::to_string_pretty(&serializable)?;
70        fs::write(path, json)?;
71        Ok(())
72    }
73
74    pub fn load(path: &Path) -> Result<Self> {
75        let json = fs::read_to_string(path)?;
76        let serializable: SerializableGraph = serde_json::from_str(&json)?;
77        
78        let mut graph = DiGraph::new();
79        for node in serializable.nodes {
80            graph.add_node(node);
81        }
82        for (from, to, edge_type) in serializable.edges {
83            graph.add_edge(NodeIndex::new(from), NodeIndex::new(to), edge_type);
84        }
85        
86        Ok(Self { graph })
87    }
88}