1use std::cell::Cell;
2use std::collections::HashSet;
3
4use rustc_hir::def_id::DefId;
5use rustc_index::IndexVec;
6use rustc_middle::mir::Local;
7use rustc_span::{DUMMY_SP, Span};
8
9pub type EdgeIdx = usize;
10pub type GraphNodes = IndexVec<Local, DataflowNode>;
11pub type GraphEdges = IndexVec<EdgeIdx, DataflowEdge>;
12
13#[derive(Clone, Debug)]
14pub enum NodeOp {
15 Nop,
16 Err,
17 Const(String, String),
18 Use,
19 Repeat,
20 Ref,
21 ThreadLocalRef,
22 AddressOf,
23 Len,
24 Cast,
25 BinaryOp,
26 CheckedBinaryOp,
27 #[cfg(not(rapx_rustc_ge_196))]
28 NullaryOp,
29 UnaryOp,
30 Discriminant,
31 Aggregate(AggKind),
32 #[cfg(not(rapx_rustc_ge_196))]
33 ShallowInitBox,
34 CopyForDeref,
35 RawPtr,
36 Call(DefId),
37 CallOperand,
38}
39
40#[derive(Clone, Debug)]
41pub enum EdgeOp {
42 Nop,
43 Move,
44 Copy,
45 Const,
46 Immut,
47 Mut,
48 Deref,
49 Field(usize),
50 Downcast(String),
51 Index,
52 ConstIndex,
53 SubSlice,
54}
55
56#[derive(Clone, Copy, Debug)]
57pub enum AggKind {
58 Array,
59 Tuple,
60 Adt(DefId),
61 Closure(DefId),
62 Coroutine(DefId),
63 RawPtr,
64}
65
66#[derive(Clone, Debug)]
67pub struct DataflowEdge {
68 pub src: Local,
69 pub dst: Local,
70 pub op: EdgeOp,
71 pub seq: usize,
72 pub block: usize,
73 pub statement_index: usize,
74}
75
76#[derive(Clone, Debug)]
77pub struct DataflowNode {
78 pub ops: Vec<NodeOp>,
79 pub span: Span,
80 pub seq: usize,
81 pub out_edges: Vec<EdgeIdx>,
82 pub in_edges: Vec<EdgeIdx>,
83}
84
85impl DataflowNode {
86 pub fn new() -> Self {
87 Self {
88 ops: vec![NodeOp::Nop],
89 span: DUMMY_SP,
90 seq: 0,
91 out_edges: vec![],
92 in_edges: vec![],
93 }
94 }
95}
96
97#[derive(Clone)]
98pub struct DataflowGraph {
99 pub def_id: DefId,
100 pub span: Span,
101 pub argc: usize,
102 pub nodes: GraphNodes,
103 pub edges: GraphEdges,
104 pub n_locals: usize,
105 pub closures: HashSet<DefId>,
106 pub block: usize,
107 pub statement_index: usize,
108}
109
110impl DataflowGraph {
111 pub fn new(def_id: DefId, span: Span, argc: usize, n_locals: usize) -> Self {
112 Self {
113 def_id,
114 span,
115 argc,
116 nodes: GraphNodes::from_elem_n(DataflowNode::new(), n_locals),
117 edges: GraphEdges::new(),
118 n_locals,
119 closures: HashSet::new(),
120 block: 0,
121 statement_index: 0,
122 }
123 }
124
125 pub fn node(&self, local: Local) -> &DataflowNode {
126 &self.nodes[local]
127 }
128
129 pub fn node_mut(&mut self, local: Local) -> &mut DataflowNode {
130 &mut self.nodes[local]
131 }
132
133 pub fn edge(&self, idx: EdgeIdx) -> &DataflowEdge {
134 &self.edges[idx]
135 }
136
137 pub fn is_marker(&self, idx: Local) -> bool {
138 idx >= Local::from_usize(self.n_locals)
139 }
140
141 pub fn add_node_edge(&mut self, src: Local, dst: Local, op: EdgeOp) -> EdgeIdx {
142 let seq = self.nodes[dst].seq;
143 let edge_idx = self.edges.push(DataflowEdge {
144 src,
145 dst,
146 op,
147 seq,
148 block: self.block,
149 statement_index: self.statement_index,
150 });
151 self.nodes[dst].in_edges.push(edge_idx);
152 self.nodes[src].out_edges.push(edge_idx);
153 edge_idx
154 }
155
156 pub fn add_const_edge(
157 &mut self,
158 src_desc: String,
159 src_ty: String,
160 dst: Local,
161 op: EdgeOp,
162 ) -> EdgeIdx {
163 let seq = self.nodes[dst].seq;
164 let mut const_node = DataflowNode::new();
165 const_node.ops[0] = NodeOp::Const(src_desc, src_ty);
166 let src = self.nodes.push(const_node);
167 let edge_idx = self.edges.push(DataflowEdge {
168 src,
169 dst,
170 op,
171 seq,
172 block: self.block,
173 statement_index: self.statement_index,
174 });
175 self.nodes[dst].in_edges.push(edge_idx);
176 edge_idx
177 }
178
179 pub fn get_upside_idx(&self, node_idx: Local, order: usize) -> Option<Local> {
180 if let Some(edge_idx) = self.nodes[node_idx].in_edges.get(order) {
181 Some(self.edges[*edge_idx].src)
182 } else {
183 None
184 }
185 }
186
187 pub fn get_downside_idx(&self, node_idx: Local, order: usize) -> Option<Local> {
188 if let Some(edge_idx) = self.nodes[node_idx].out_edges.get(order) {
189 Some(self.edges[*edge_idx].dst)
190 } else {
191 None
192 }
193 }
194
195 pub fn is_connected(&self, idx_1: Local, idx_2: Local) -> bool {
196 let target = idx_2;
197 let find = Cell::new(false);
198 let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
199 find.set(idx == target);
200 if find.get() {
201 DFSStatus::Stop
202 } else {
203 DFSStatus::Continue
204 }
205 };
206 let mut seen = HashSet::new();
207 self.dfs(
208 idx_1,
209 Direction::Downside,
210 &mut node_operator,
211 &mut Self::always_true_edge_validator,
212 false,
213 &mut seen,
214 );
215 seen.clear();
216 if !find.get() {
217 self.dfs(
218 idx_1,
219 Direction::Upside,
220 &mut node_operator,
221 &mut Self::always_true_edge_validator,
222 false,
223 &mut seen,
224 );
225 }
226 find.get()
227 }
228
229 pub fn param_return_deps(&self) -> IndexVec<Local, bool> {
230 let _0 = Local::from_usize(0);
231 let deps = (0..self.argc + 1)
232 .map(|i| {
233 let _i = Local::from_usize(i);
234 self.is_connected(_i, _0)
235 })
236 .collect();
237 deps
238 }
239
240 pub fn dfs<F, G>(
241 &self,
242 now: Local,
243 direction: Direction,
244 node_operator: &mut F,
245 edge_validator: &mut G,
246 traverse_all: bool,
247 seen: &mut HashSet<Local>,
248 ) -> (DFSStatus, bool)
249 where
250 F: FnMut(&DataflowGraph, Local) -> DFSStatus,
251 G: FnMut(&DataflowGraph, EdgeIdx) -> DFSStatus,
252 {
253 if seen.contains(&now) {
254 return (DFSStatus::Stop, false);
255 }
256 seen.insert(now);
257 macro_rules! traverse {
258 ($edges: ident, $field: ident) => {
259 for edge_idx in self.nodes[now].$edges.iter() {
260 let edge = &self.edges[*edge_idx];
261 if matches!(edge_validator(self, *edge_idx), DFSStatus::Continue) {
262 let (dfs_status, result) = self.dfs(
263 edge.$field,
264 direction,
265 node_operator,
266 edge_validator,
267 traverse_all,
268 seen,
269 );
270 if matches!(dfs_status, DFSStatus::Stop) && result && !traverse_all {
271 return (DFSStatus::Stop, true);
272 }
273 }
274 }
275 };
276 }
277 if matches!(node_operator(self, now), DFSStatus::Continue) {
278 match direction {
279 Direction::Upside => {
280 traverse!(in_edges, src);
281 }
282 Direction::Downside => {
283 traverse!(out_edges, dst);
284 }
285 Direction::Both => {
286 traverse!(in_edges, src);
287 traverse!(out_edges, dst);
288 }
289 };
290 (DFSStatus::Continue, false)
291 } else {
292 (DFSStatus::Stop, true)
293 }
294 }
295
296 pub fn equivalent_edge_validator(graph: &DataflowGraph, idx: EdgeIdx) -> DFSStatus {
297 match graph.edges[idx].op {
298 EdgeOp::Copy | EdgeOp::Move | EdgeOp::Mut | EdgeOp::Immut | EdgeOp::Deref => {
299 DFSStatus::Continue
300 }
301 EdgeOp::Nop
302 | EdgeOp::Const
303 | EdgeOp::Downcast(_)
304 | EdgeOp::Field(_)
305 | EdgeOp::Index
306 | EdgeOp::ConstIndex
307 | EdgeOp::SubSlice => DFSStatus::Stop,
308 }
309 }
310
311 pub fn always_true_edge_validator(_: &DataflowGraph, _: EdgeIdx) -> DFSStatus {
312 DFSStatus::Continue
313 }
314
315 pub fn collect_equivalent_locals(&self, local: Local, strict: bool) -> HashSet<Local> {
316 let mut set = HashSet::new();
317 let root = Cell::new(local);
318 let reduce_func = if strict {
319 DFSStatus::and
320 } else {
321 DFSStatus::or
322 };
323 let mut find_root_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
324 let node = &graph.nodes[idx];
325 node.ops
326 .iter()
327 .map(|op| match op {
328 NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
329 root.set(idx);
330 DFSStatus::Continue
331 }
332 NodeOp::Call(_) => {
333 root.set(idx);
334 DFSStatus::Stop
335 }
336 _ => DFSStatus::Stop,
337 })
338 .reduce(reduce_func)
339 .unwrap()
340 };
341 let mut find_equivalent_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
342 let node = &graph.nodes[idx];
343 if set.contains(&idx) {
344 return DFSStatus::Stop;
345 }
346 node.ops
347 .iter()
348 .map(|op| match op {
349 NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
350 set.insert(idx);
351 DFSStatus::Continue
352 }
353 NodeOp::Call(_) => {
354 if idx == root.get() {
355 set.insert(idx);
356 DFSStatus::Continue
357 } else {
358 DFSStatus::Stop
359 }
360 }
361 _ => DFSStatus::Stop,
362 })
363 .reduce(reduce_func)
364 .unwrap()
365 };
366 let mut seen = HashSet::new();
367 self.dfs(
368 local,
369 Direction::Upside,
370 &mut find_root_operator,
371 &mut Self::equivalent_edge_validator,
372 true,
373 &mut seen,
374 );
375 seen.clear();
376 self.dfs(
377 root.get(),
378 Direction::Downside,
379 &mut find_equivalent_operator,
380 &mut Self::equivalent_edge_validator,
381 true,
382 &mut seen,
383 );
384 set
385 }
386
387 pub fn collect_ancestor_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
388 let mut ret = HashSet::new();
389 let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
390 ret.insert(idx);
391 DFSStatus::Continue
392 };
393 let mut seen = HashSet::new();
394 self.dfs(
395 local,
396 Direction::Upside,
397 &mut node_operator,
398 &mut DataflowGraph::always_true_edge_validator,
399 true,
400 &mut seen,
401 );
402 if !self_included {
403 ret.remove(&local);
404 }
405 ret
406 }
407
408 pub fn collect_descending_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
409 let mut ret = HashSet::new();
410 let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
411 ret.insert(idx);
412 DFSStatus::Continue
413 };
414 let mut seen = HashSet::new();
415 self.dfs(
416 local,
417 Direction::Downside,
418 &mut node_operator,
419 &mut DataflowGraph::always_true_edge_validator,
420 true,
421 &mut seen,
422 );
423 if !self_included {
424 ret.remove(&local);
425 }
426 ret
427 }
428
429 pub fn get_field_sequence(&self, local: Local) -> Option<(Local, Vec<usize>)> {
430 let mut fields = vec![];
431 let var = Cell::new(local);
432 let mut node_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
433 if graph.is_marker(idx) {
434 DFSStatus::Continue
435 } else {
436 var.set(idx);
437 DFSStatus::Stop
438 }
439 };
440 let mut edge_validator = |graph: &DataflowGraph, idx: EdgeIdx| -> DFSStatus {
441 if let EdgeOp::Field(field) = graph.edges[idx].op {
442 fields.insert(0, field);
443 DFSStatus::Continue
444 } else {
445 DFSStatus::Stop
446 }
447 };
448 let mut seen = HashSet::new();
449 self.dfs(
450 local,
451 Direction::Upside,
452 &mut node_operator,
453 &mut edge_validator,
454 false,
455 &mut seen,
456 );
457 if fields.is_empty() {
458 None
459 } else {
460 Some((var.get(), fields))
461 }
462 }
463}
464
465#[derive(Clone, Copy)]
466pub enum Direction {
467 Upside,
468 Downside,
469 Both,
470}
471
472pub enum DFSStatus {
473 Continue,
474 Stop,
475}
476
477impl DFSStatus {
478 pub fn and(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
479 if matches!(s1, DFSStatus::Stop) || matches!(s2, DFSStatus::Stop) {
480 DFSStatus::Stop
481 } else {
482 DFSStatus::Continue
483 }
484 }
485
486 pub fn or(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
487 if matches!(s1, DFSStatus::Continue) || matches!(s2, DFSStatus::Continue) {
488 DFSStatus::Continue
489 } else {
490 DFSStatus::Stop
491 }
492 }
493}