rapx/analysis/safetyflow_analysis/
safetyflow_graph.rs

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

rapx/analysis/safetyflow_analysis/
safetyflow_graph.rs
