Skip to main content

icb_graph/
graph.rs

1use icb_common::NodeKind;
2use petgraph::stable_graph::StableGraph;
3use serde::{Deserialize, Serialize};
4
5/// A node in the Code Property Graph.
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct Node {
8    pub kind: NodeKind,
9    pub name: Option<String>,
10    pub usr: Option<String>,
11    pub start_line: usize,
12    pub end_line: usize,
13}
14
15/// Edge type.
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17pub enum Edge {
18    AstChild,
19    Call,
20    Reference,
21}
22
23pub type CpgGraph = StableGraph<Node, Edge>;
24
25/// Central Code Property Graph container.
26#[derive(Debug)]
27pub struct CodePropertyGraph {
28    pub graph: CpgGraph,
29}
30
31impl CodePropertyGraph {
32    pub fn new() -> Self {
33        Self {
34            graph: StableGraph::new(),
35        }
36    }
37
38    pub fn node_count(&self) -> usize {
39        self.graph.node_count()
40    }
41
42    pub fn edge_count(&self) -> usize {
43        self.graph.edge_count()
44    }
45}
46
47impl Default for CodePropertyGraph {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53/// Intermediate representation for (de)serialization.
54#[derive(Serialize, Deserialize)]
55pub struct GraphData {
56    pub nodes: Vec<Node>,
57    pub edges: Vec<(usize, usize, Edge)>,
58}
59
60impl From<&CodePropertyGraph> for GraphData {
61    fn from(cpg: &CodePropertyGraph) -> Self {
62        let nodes: Vec<Node> = cpg.graph.node_weights().cloned().collect();
63        let mut edges = Vec::new();
64        for edge_idx in cpg.graph.edge_indices() {
65            if let Some((src, tgt)) = cpg.graph.edge_endpoints(edge_idx) {
66                let weight = cpg.graph[edge_idx].clone();
67                edges.push((src.index(), tgt.index(), weight));
68            }
69        }
70        GraphData { nodes, edges }
71    }
72}
73
74impl From<GraphData> for CodePropertyGraph {
75    fn from(data: GraphData) -> Self {
76        let mut graph: CpgGraph = StableGraph::new();
77        let mut indices: Vec<petgraph::stable_graph::NodeIndex> =
78            Vec::with_capacity(data.nodes.len());
79        for node in data.nodes {
80            let idx = graph.add_node(node);
81            indices.push(idx);
82        }
83        for (src, tgt, edge) in data.edges {
84            if src < indices.len() && tgt < indices.len() {
85                graph.add_edge(indices[src], indices[tgt], edge);
86            }
87        }
88        CodePropertyGraph { graph }
89    }
90}