rapx/analysis/upg/
upg_graph.rs

1use std::{
2    collections::{HashMap, HashSet},
3    fmt::{self, Write},
4};
5
6use crate::{analysis::utils::fn_info::*, utils::source::get_adt_name};
7use rustc_hir::{Safety, def_id::DefId};
8use rustc_middle::ty::TyCtxt;
9
10use super::upg_unit::UPGUnit;
11use petgraph::{
12    Graph,
13    dot::{Config, Dot},
14    graph::{DiGraph, EdgeReference, NodeIndex},
15};
16
17#[derive(Debug, Clone, Eq, PartialEq, Hash)]
18pub enum UPGNode {
19    SafeFn(DefId, String),
20    UnsafeFn(DefId, String),
21    MergedCallerCons(String),
22    MutMethods(String),
23}
24
25#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
26pub enum UPGEdge {
27    CallerToCallee,
28    ConsToMethod,
29    MutToCaller,
30}
31
32impl UPGNode {
33    pub fn from(node: FnInfo) -> Self {
34        match node.fn_safety {
35            Safety::Unsafe => match node.fn_kind {
36                FnKind::Constructor => UPGNode::UnsafeFn(node.def_id, "box".to_string()),
37                FnKind::Method => UPGNode::UnsafeFn(node.def_id, "ellipse".to_string()),
38                FnKind::Fn => UPGNode::UnsafeFn(node.def_id, "box".to_string()),
39                FnKind::Intrinsic => UPGNode::UnsafeFn(node.def_id, "box".to_string()),
40            },
41            Safety::Safe => match node.fn_kind {
42                FnKind::Constructor => UPGNode::SafeFn(node.def_id, "box".to_string()),
43                FnKind::Method => UPGNode::SafeFn(node.def_id, "ellipse".to_string()),
44                FnKind::Fn => UPGNode::SafeFn(node.def_id, "box".to_string()),
45                FnKind::Intrinsic => UPGNode::SafeFn(node.def_id, "box".to_string()),
46            },
47        }
48    }
49}
50
51impl fmt::Display for UPGNode {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        match self {
54            UPGNode::SafeFn(_, _) => write!(f, "Safe"),
55            UPGNode::UnsafeFn(_, _) => write!(f, "Unsafe"),
56            UPGNode::MergedCallerCons(_) => write!(f, "MergedCallerCons"),
57            UPGNode::MutMethods(_) => write!(f, "MutMethods"),
58        }
59    }
60}
61
62impl fmt::Display for UPGEdge {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            UPGEdge::CallerToCallee => write!(f, "CallerToCallee"),
66            UPGEdge::ConsToMethod => write!(f, "ConsToMethod"),
67            UPGEdge::MutToCaller => write!(f, "MutToCaller"),
68        }
69    }
70}
71
72/// Holds graph data for a single module before DOT generation.
73pub struct UPGraph {
74    // Nodes grouped by their associated struct/type name.
75    structs: HashMap<String, HashSet<FnInfo>>,
76    // Edges between DefIds with their type.
77    edges: HashSet<(DefId, DefId, UPGEdge)>,
78    // Pre-generated DOT attribute strings for each node (DefId).
79    nodes: HashMap<DefId, String>,
80}
81
82impl UPGraph {
83    pub fn new() -> Self {
84        Self {
85            structs: HashMap::new(),
86            edges: HashSet::new(),
87            nodes: HashMap::new(),
88        }
89    }
90
91    pub fn add_node(&mut self, tcx: TyCtxt<'_>, node: FnInfo, custom_label: Option<String>) {
92        let adt_name = get_adt_name(tcx, node.def_id);
93        self.structs.entry(adt_name).or_default().insert(node);
94
95        if !self.nodes.contains_key(&node.def_id) || custom_label.is_some() {
96            let attr = if let Some(label) = custom_label {
97                if node.fn_kind == FnKind::Constructor {
98                    format!(
99                        "label=\"{}\", shape=\"septagon\", style=\"filled\", fillcolor=\"#f0f0f0\", color=\"#555555\"",
100                        label
101                    )
102                } else {
103                    format!(
104                        "label=\"{}\", shape=\"ellipse\", style=\"filled\", fillcolor=\"#f0f0f0\", color=\"#555555\"",
105                        label
106                    )
107                }
108            } else {
109                let upg_node = UPGNode::from(node);
110                Self::node_to_dot_attr(&upg_node)
111            };
112
113            self.nodes.insert(node.def_id, attr);
114        }
115    }
116
117    pub fn add_edge(&mut self, from: DefId, to: DefId, edge_type: UPGEdge) {
118        if from == to {
119            return;
120        }
121        self.edges.insert((from, to, edge_type));
122    }
123
124    pub fn upg_unit_string(&self, module_name: &str) -> String {
125        let mut dot = String::new();
126        let graph_id = module_name
127            .replace("::", "_")
128            .replace(|c: char| !c.is_alphanumeric(), "_");
129
130        writeln!(dot, "digraph {} {{", graph_id).unwrap();
131        writeln!(dot, "    compound=true;").unwrap();
132        writeln!(dot, "    rankdir=LR;").unwrap();
133
134        for (struct_name, nodes) in &self.structs {
135            let cluster_id = format!(
136                "cluster_{}",
137                struct_name.replace(|c: char| !c.is_alphanumeric(), "_")
138            );
139
140            writeln!(dot, "    subgraph {} {{", cluster_id).unwrap();
141            writeln!(dot, "        label=\"{}\";", struct_name).unwrap();
142            writeln!(dot, "        style=dashed;").unwrap();
143            writeln!(dot, "        color=gray;").unwrap();
144
145            for node in nodes {
146                let def_id = node.def_id;
147                let node_id =
148                    format!("n_{:?}", def_id).replace(|c: char| !c.is_alphanumeric(), "_");
149
150                if let Some(attr) = self.nodes.get(&def_id) {
151                    writeln!(dot, "        {} [{}];", node_id, attr).unwrap();
152                }
153            }
154            writeln!(dot, "    }}").unwrap();
155        }
156
157        for (from, to, edge_type) in &self.edges {
158            let from_id = format!("n_{:?}", from).replace(|c: char| !c.is_alphanumeric(), "_");
159            let to_id = format!("n_{:?}", to).replace(|c: char| !c.is_alphanumeric(), "_");
160
161            let attr = match edge_type {
162                UPGEdge::CallerToCallee => "color=black, style=solid",
163                UPGEdge::ConsToMethod => "color=black, style=dotted",
164                UPGEdge::MutToCaller => "color=blue, style=dashed",
165            };
166
167            writeln!(dot, "    {} -> {} [{}];", from_id, to_id, attr).unwrap();
168        }
169
170        writeln!(dot, "}}").unwrap();
171        dot
172    }
173
174    fn node_to_dot_attr(node: &UPGNode) -> String {
175        match node {
176            UPGNode::SafeFn(def_id, shape) => {
177                format!("label=\"{:?}\", color=black, shape={:?}", def_id, shape)
178            }
179            UPGNode::UnsafeFn(def_id, shape) => {
180                let label = format!("{:?}", def_id);
181                format!("label=\"{}\", shape={:?}, color=red", label, shape)
182            }
183            _ => "label=\"Unknown\"".to_string(),
184        }
185    }
186
187    pub fn generate_dot_from_upg_unit(upg: &UPGUnit) -> String {
188        let mut graph: Graph<UPGNode, UPGEdge> = DiGraph::new();
189        let get_edge_attr = |_graph: &Graph<UPGNode, UPGEdge>,
190                             edge_ref: EdgeReference<'_, UPGEdge>| {
191            match edge_ref.weight() {
192                UPGEdge::CallerToCallee => "color=black, style=solid",
193                UPGEdge::ConsToMethod => "color=black, style=dotted",
194                UPGEdge::MutToCaller => "color=blue, style=dashed",
195            }
196            .to_owned()
197        };
198        let get_node_attr =
199            |_graph: &Graph<UPGNode, UPGEdge>, node_ref: (NodeIndex, &UPGNode)| match node_ref.1 {
200                UPGNode::SafeFn(def_id, shape) => {
201                    format!("label=\"{:?}\", color=black, shape={:?}", def_id, shape)
202                }
203                UPGNode::UnsafeFn(def_id, shape) => {
204                    let label = format!("{:?}\n ", def_id);
205                    let node_attr = format!("label={:?}, shape={:?}, color=red", label, shape);
206                    node_attr
207                }
208                UPGNode::MergedCallerCons(label) => {
209                    format!(
210                        "label=\"{}\", shape=box, style=filled, fillcolor=lightgrey",
211                        label
212                    )
213                }
214                UPGNode::MutMethods(label) => {
215                    format!(
216                        "label=\"{}\", shape=octagon, style=filled, fillcolor=lightyellow",
217                        label
218                    )
219                }
220            };
221
222        let caller_node = graph.add_node(UPGNode::from(upg.caller));
223        if !upg.caller_cons.is_empty() {
224            let cons_labels: Vec<String> = upg
225                .caller_cons
226                .iter()
227                .map(|con| format!("{:?}", con.def_id))
228                .collect();
229            let merged_label = format!("Caller Constructors\n{}", cons_labels.join("\n"));
230            let merged_cons_node = graph.add_node(UPGNode::MergedCallerCons(merged_label));
231            graph.add_edge(merged_cons_node, caller_node, UPGEdge::ConsToMethod);
232        }
233
234        if !upg.mut_methods.is_empty() {
235            let mut_method_labels: Vec<String> = upg
236                .mut_methods
237                .iter()
238                .map(|def_id| format!("{:?}", def_id))
239                .collect();
240            let merged_label = format!("Mutable Methods\n{}", mut_method_labels.join("\n"));
241
242            let mut_methods_node = graph.add_node(UPGNode::MutMethods(merged_label));
243            graph.add_edge(mut_methods_node, caller_node, UPGEdge::MutToCaller);
244        }
245
246        /*
247        for (callee, cons) in &self.callee_cons_pair {
248            let callee_node = graph.add_node(Self::get_node(*callee));
249            for callee_cons in cons {
250                let callee_cons_node = graph.add_node(Self::get_node(*callee_cons));
251                graph.add_edge(callee_cons_node, callee_node, UPGEdge::ConsToMethod);
252            }
253            graph.add_edge(caller_node, callee_node, UPGEdge::CallerToCallee);
254        }
255        */
256        let mut dot_str = String::new();
257        let dot = Dot::with_attr_getters(
258            &graph,
259            // &[Config::NodeNoLabel, Config::EdgeNoLabel],
260            &[Config::NodeNoLabel],
261            &get_edge_attr,
262            &get_node_attr,
263        );
264
265        write!(dot_str, "{}", dot).unwrap();
266        // println!("{}", dot_str);
267        dot_str
268    }
269}
rapx/analysis/upg/upg_graph.rs

rapx/analysis/upg/
upg_graph.rs
