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    pub docstring: Option<String>,
23}
24
25#[derive(Serialize, Deserialize, Debug, Clone)]
26pub enum EdgeType {
27    Contains,
28    Imports,
29    Calls,
30}
31
32#[derive(Serialize, Deserialize)]
33struct SerializableGraph {
34    nodes: Vec<NodeData>,
35    edges: Vec<(usize, usize, EdgeType)>,
36}
37
38pub struct ProjectGraph {
39    pub graph: DiGraph<NodeData, EdgeType>,
40}
41
42impl ProjectGraph {
43    pub fn new() -> Self {
44        Self {
45            graph: DiGraph::new(),
46        }
47    }
48
49    pub fn add_node(&mut self, data: NodeData) -> NodeIndex {
50        self.graph.add_node(data)
51    }
52
53    pub fn add_edge(&mut self, from: NodeIndex, to: NodeIndex, edge_type: EdgeType) {
54        self.graph.add_edge(from, to, edge_type);
55    }
56
57    pub fn save(&self, path: &Path) -> Result<()> {
58        let mut nodes = Vec::new();
59        for i in 0..self.graph.node_count() {
60            nodes.push(self.graph[NodeIndex::new(i)].clone());
61        }
62
63        let mut edges = Vec::new();
64        for edge in self.graph.edge_indices() {
65            let (from, to) = self.graph.edge_endpoints(edge).unwrap();
66            edges.push((from.index(), to.index(), self.graph[edge].clone()));
67        }
68
69        let serializable = SerializableGraph { nodes, edges };
70        let json = serde_json::to_string_pretty(&serializable)?;
71        fs::write(path, json)?;
72        Ok(())
73    }
74
75    pub fn load(path: &Path) -> Result<Self> {
76        let json = fs::read_to_string(path)?;
77        let serializable: SerializableGraph = serde_json::from_str(&json)?;
78        
79        let mut graph = DiGraph::new();
80        for node in serializable.nodes {
81            graph.add_node(node);
82        }
83        for (from, to, edge_type) in serializable.edges {
84            graph.add_edge(NodeIndex::new(from), NodeIndex::new(to), edge_type);
85        }
86        
87        Ok(Self { graph })
88    }
89}
project_map_cli_rust/core/graph.rs

project_map_cli_rust/core/
graph.rs
