rapx/analysis/dataflow/
types.rs

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_ge_99))]
28    NullaryOp,
29    UnaryOp,
30    Discriminant,
31    Aggregate(AggKind),
32    #[cfg(not(rapx_ge_99))]
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}
107
108impl DataflowGraph {
109    pub fn new(def_id: DefId, span: Span, argc: usize, n_locals: usize) -> Self {
110        Self {
111            def_id,
112            span,
113            argc,
114            nodes: GraphNodes::from_elem_n(DataflowNode::new(), n_locals),
115            edges: GraphEdges::new(),
116            n_locals,
117            closures: HashSet::new(),
118        }
119    }
120
121    pub fn node(&self, local: Local) -> &DataflowNode {
122        &self.nodes[local]
123    }
124
125    pub fn node_mut(&mut self, local: Local) -> &mut DataflowNode {
126        &mut self.nodes[local]
127    }
128
129    pub fn edge(&self, idx: EdgeIdx) -> &DataflowEdge {
130        &self.edges[idx]
131    }
132
133    pub fn is_marker(&self, idx: Local) -> bool {
134        idx >= Local::from_usize(self.n_locals)
135    }
136
137    pub fn add_node_edge(
138        &mut self,
139        src: Local,
140        dst: Local,
141        op: EdgeOp,
142        block: usize,
143        statement_index: usize,
144    ) -> EdgeIdx {
145        let seq = self.nodes[dst].seq;
146        let edge_idx = self.edges.push(DataflowEdge {
147            src,
148            dst,
149            op,
150            seq,
151            block,
152            statement_index,
153        });
154        self.nodes[dst].in_edges.push(edge_idx);
155        self.nodes[src].out_edges.push(edge_idx);
156        edge_idx
157    }
158
159    pub fn add_const_edge(
160        &mut self,
161        src_desc: String,
162        src_ty: String,
163        dst: Local,
164        op: EdgeOp,
165        block: usize,
166        statement_index: usize,
167    ) -> EdgeIdx {
168        let seq = self.nodes[dst].seq;
169        let mut const_node = DataflowNode::new();
170        const_node.ops[0] = NodeOp::Const(src_desc, src_ty);
171        let src = self.nodes.push(const_node);
172        let edge_idx = self.edges.push(DataflowEdge {
173            src,
174            dst,
175            op,
176            seq,
177            block,
178            statement_index,
179        });
180        self.nodes[dst].in_edges.push(edge_idx);
181        edge_idx
182    }
183
184    pub fn get_upside_idx(&self, node_idx: Local, order: usize) -> Option<Local> {
185        if let Some(edge_idx) = self.nodes[node_idx].in_edges.get(order) {
186            Some(self.edges[*edge_idx].src)
187        } else {
188            None
189        }
190    }
191
192    pub fn get_downside_idx(&self, node_idx: Local, order: usize) -> Option<Local> {
193        if let Some(edge_idx) = self.nodes[node_idx].out_edges.get(order) {
194            Some(self.edges[*edge_idx].dst)
195        } else {
196            None
197        }
198    }
199
200    pub fn is_connected(&self, idx_1: Local, idx_2: Local) -> bool {
201        let target = idx_2;
202        let find = Cell::new(false);
203        let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
204            find.set(idx == target);
205            if find.get() {
206                DFSStatus::Stop
207            } else {
208                DFSStatus::Continue
209            }
210        };
211        let mut seen = HashSet::new();
212        self.dfs(
213            idx_1,
214            Direction::Downside,
215            &mut node_operator,
216            &mut Self::always_true_edge_validator,
217            false,
218            &mut seen,
219        );
220        seen.clear();
221        if !find.get() {
222            self.dfs(
223                idx_1,
224                Direction::Upside,
225                &mut node_operator,
226                &mut Self::always_true_edge_validator,
227                false,
228                &mut seen,
229            );
230        }
231        find.get()
232    }
233
234    pub fn param_return_deps(&self) -> IndexVec<Local, bool> {
235        let ret_local = Local::from_usize(0);
236        let deps = (0..self.argc + 1)
237            .map(|i| {
238                let arg_local = Local::from_usize(i);
239                self.is_connected(arg_local, ret_local)
240            })
241            .collect();
242        deps
243    }
244
245    pub fn dfs<F, G>(
246        &self,
247        now: Local,
248        direction: Direction,
249        node_operator: &mut F,
250        edge_validator: &mut G,
251        traverse_all: bool,
252        seen: &mut HashSet<Local>,
253    ) -> (DFSStatus, bool)
254    where
255        F: FnMut(&DataflowGraph, Local) -> DFSStatus,
256        G: FnMut(&DataflowGraph, EdgeIdx) -> DFSStatus,
257    {
258        if seen.contains(&now) {
259            return (DFSStatus::Stop, false);
260        }
261        seen.insert(now);
262        macro_rules! traverse {
263            ($edges: ident, $field: ident) => {
264                for edge_idx in self.nodes[now].$edges.iter() {
265                    let edge = &self.edges[*edge_idx];
266                    if matches!(edge_validator(self, *edge_idx), DFSStatus::Continue) {
267                        let (dfs_status, result) = self.dfs(
268                            edge.$field,
269                            direction,
270                            node_operator,
271                            edge_validator,
272                            traverse_all,
273                            seen,
274                        );
275                        if matches!(dfs_status, DFSStatus::Stop) && result && !traverse_all {
276                            return (DFSStatus::Stop, true);
277                        }
278                    }
279                }
280            };
281        }
282        if matches!(node_operator(self, now), DFSStatus::Continue) {
283            match direction {
284                Direction::Upside => {
285                    traverse!(in_edges, src);
286                }
287                Direction::Downside => {
288                    traverse!(out_edges, dst);
289                }
290                Direction::Both => {
291                    traverse!(in_edges, src);
292                    traverse!(out_edges, dst);
293                }
294            };
295            (DFSStatus::Continue, false)
296        } else {
297            (DFSStatus::Stop, true)
298        }
299    }
300
301    pub fn find_first_node<P, E>(
302        &self,
303        start: Local,
304        direction: Direction,
305        node_predicate: &mut P,
306        edge_validator: &mut E,
307    ) -> Option<Local>
308    where
309        P: FnMut(&DataflowGraph, Local) -> bool,
310        E: FnMut(&DataflowGraph, EdgeIdx) -> DFSStatus,
311    {
312        let mut result = None;
313        let mut node_op = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
314            if node_predicate(graph, idx) {
315                result = Some(idx);
316                DFSStatus::Stop
317            } else {
318                DFSStatus::Continue
319            }
320        };
321        let mut seen = HashSet::new();
322        self.dfs(
323            start,
324            direction,
325            &mut node_op,
326            edge_validator,
327            false,
328            &mut seen,
329        );
330        result
331    }
332
333    pub fn find_all_nodes<P, E>(
334        &self,
335        start: Local,
336        direction: Direction,
337        node_predicate: &mut P,
338        edge_validator: &mut E,
339    ) -> Vec<Local>
340    where
341        P: FnMut(&DataflowGraph, Local) -> bool,
342        E: FnMut(&DataflowGraph, EdgeIdx) -> DFSStatus,
343    {
344        let mut results = Vec::new();
345        let mut node_op = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
346            if node_predicate(graph, idx) {
347                results.push(idx);
348            }
349            DFSStatus::Continue
350        };
351        let mut seen = HashSet::new();
352        self.dfs(
353            start,
354            direction,
355            &mut node_op,
356            edge_validator,
357            true,
358            &mut seen,
359        );
360        results
361    }
362
363    pub fn equivalent_edge_validator(graph: &DataflowGraph, idx: EdgeIdx) -> DFSStatus {
364        match graph.edges[idx].op {
365            EdgeOp::Copy | EdgeOp::Move | EdgeOp::Mut | EdgeOp::Immut | EdgeOp::Deref => {
366                DFSStatus::Continue
367            }
368            EdgeOp::Nop
369            | EdgeOp::Const
370            | EdgeOp::Downcast(_)
371            | EdgeOp::Field(_)
372            | EdgeOp::Index
373            | EdgeOp::ConstIndex
374            | EdgeOp::SubSlice => DFSStatus::Stop,
375        }
376    }
377
378    pub fn always_true_edge_validator(_: &DataflowGraph, _: EdgeIdx) -> DFSStatus {
379        DFSStatus::Continue
380    }
381
382    pub fn collect_equivalent_locals(&self, local: Local, strict: bool) -> HashSet<Local> {
383        let mut set = HashSet::new();
384        let root = Cell::new(local);
385        let reduce_func = if strict {
386            DFSStatus::and
387        } else {
388            DFSStatus::or
389        };
390        let mut find_root_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
391            let node = &graph.nodes[idx];
392            node.ops
393                .iter()
394                .map(|op| match op {
395                    NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
396                        root.set(idx);
397                        DFSStatus::Continue
398                    }
399                    NodeOp::Call(_) => {
400                        root.set(idx);
401                        DFSStatus::Stop
402                    }
403                    _ => DFSStatus::Stop,
404                })
405                .reduce(reduce_func)
406                .unwrap()
407        };
408        let mut find_equivalent_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
409            let node = &graph.nodes[idx];
410            if set.contains(&idx) {
411                return DFSStatus::Stop;
412            }
413            node.ops
414                .iter()
415                .map(|op| match op {
416                    NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
417                        set.insert(idx);
418                        DFSStatus::Continue
419                    }
420                    NodeOp::Call(_) => {
421                        if idx == root.get() {
422                            set.insert(idx);
423                            DFSStatus::Continue
424                        } else {
425                            DFSStatus::Stop
426                        }
427                    }
428                    _ => DFSStatus::Stop,
429                })
430                .reduce(reduce_func)
431                .unwrap()
432        };
433        let mut seen = HashSet::new();
434        self.dfs(
435            local,
436            Direction::Upside,
437            &mut find_root_operator,
438            &mut Self::equivalent_edge_validator,
439            true,
440            &mut seen,
441        );
442        seen.clear();
443        self.dfs(
444            root.get(),
445            Direction::Downside,
446            &mut find_equivalent_operator,
447            &mut Self::equivalent_edge_validator,
448            true,
449            &mut seen,
450        );
451        set
452    }
453
454    fn collect_by_direction(
455        &self,
456        local: Local,
457        self_included: bool,
458        direction: Direction,
459    ) -> HashSet<Local> {
460        let mut ret = HashSet::new();
461        let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
462            ret.insert(idx);
463            DFSStatus::Continue
464        };
465        let mut seen = HashSet::new();
466        self.dfs(
467            local,
468            direction,
469            &mut node_operator,
470            &mut DataflowGraph::always_true_edge_validator,
471            true,
472            &mut seen,
473        );
474        if !self_included {
475            ret.remove(&local);
476        }
477        ret
478    }
479
480    pub fn collect_ancestor_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
481        self.collect_by_direction(local, self_included, Direction::Upside)
482    }
483
484    pub fn collect_descending_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
485        self.collect_by_direction(local, self_included, Direction::Downside)
486    }
487
488    pub fn get_field_sequence(&self, local: Local) -> Option<(Local, Vec<usize>)> {
489        let mut fields = vec![];
490        let var = Cell::new(local);
491        let mut node_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
492            if graph.is_marker(idx) {
493                DFSStatus::Continue
494            } else {
495                var.set(idx);
496                DFSStatus::Stop
497            }
498        };
499        let mut edge_validator = |graph: &DataflowGraph, idx: EdgeIdx| -> DFSStatus {
500            if let EdgeOp::Field(field) = graph.edges[idx].op {
501                fields.insert(0, field);
502                DFSStatus::Continue
503            } else {
504                DFSStatus::Stop
505            }
506        };
507        let mut seen = HashSet::new();
508        self.dfs(
509            local,
510            Direction::Upside,
511            &mut node_operator,
512            &mut edge_validator,
513            false,
514            &mut seen,
515        );
516        if fields.is_empty() {
517            None
518        } else {
519            Some((var.get(), fields))
520        }
521    }
522
523    /// Follow Copy/Move edges upward to find the root real local behind any
524    /// copy chains (no projections). Stops at 16 hops to bound cycles.
525    pub fn trace_origin(&self, local: Local) -> Local {
526        let mut current = local;
527        let mut seen = HashSet::new();
528        for _ in 0..16 {
529            if !seen.insert(current) {
530                break;
531            }
532            let next = self.nodes[current].in_edges.iter().find_map(|&ei| {
533                let e = &self.edges[ei];
534                if matches!(e.op, EdgeOp::Copy | EdgeOp::Move) && !self.is_marker(e.src) {
535                    Some(e.src)
536                } else {
537                    None
538                }
539            });
540            match next {
541                Some(src) if src != current => current = src,
542                _ => break,
543            }
544        }
545        current
546    }
547
548    /// Return `true` when `local` originates from a tuple field destructuring
549    /// (e.g. `(tuple.0, tuple.1)` after a call returning a tuple).
550    /// Follows Copy/Move chains upward through marker nodes and checks for
551    /// any `Field` edge along the projection chain.
552    pub fn is_from_tuple_field(&self, local: Local) -> bool {
553        let mut current = local;
554        let mut seen = HashSet::new();
555        for _ in 0..8 {
556            if !seen.insert(current) {
557                return false;
558            }
559            let mut next_local = None;
560            for &ei in &self.nodes[current].in_edges {
561                let e = &self.edges[ei];
562                if !matches!(e.op, EdgeOp::Copy | EdgeOp::Move) {
563                    continue;
564                }
565                if self.is_marker(e.src) {
566                    if self.marker_chain_has_field(e.src) {
567                        return true;
568                    }
569                    if let Some(real) = self.marker_to_real(e.src) {
570                        next_local = Some(real);
571                        break;
572                    }
573                } else {
574                    next_local = Some(e.src);
575                    break;
576                }
577            }
578            match next_local {
579                Some(src) if src != current => current = src,
580                _ => return false,
581            }
582        }
583        false
584    }
585
586    /// Walk up a projection-marker chain to find the underlying real local.
587    fn marker_to_real(&self, marker: Local) -> Option<Local> {
588        let mut current = marker;
589        for _ in 0..8 {
590            if !self.is_marker(current) {
591                return Some(current);
592            }
593            current = self.nodes[current]
594                .in_edges
595                .first()
596                .map(|&ei| self.edges[ei].src)?;
597        }
598        None
599    }
600
601    /// Check whether a projection-marker chain contains a `Field` edge.
602    fn marker_chain_has_field(&self, marker: Local) -> bool {
603        let mut current = marker;
604        for _ in 0..8 {
605            if !self.is_marker(current) {
606                return false;
607            }
608            let ei = match self.nodes[current].in_edges.first() {
609                Some(&ei) => ei,
610                None => return false,
611            };
612            if matches!(self.edges[ei].op, EdgeOp::Field(_)) {
613                return true;
614            }
615            current = self.edges[ei].src;
616        }
617        false
618    }
619}
620
621#[derive(Clone, Copy)]
622pub enum Direction {
623    Upside,
624    Downside,
625    Both,
626}
627
628pub enum DFSStatus {
629    Continue,
630    Stop,
631}
632
633impl DFSStatus {
634    pub fn and(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
635        if matches!(s1, DFSStatus::Stop) || matches!(s2, DFSStatus::Stop) {
636            DFSStatus::Stop
637        } else {
638            DFSStatus::Continue
639        }
640    }
641
642    pub fn or(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
643        if matches!(s1, DFSStatus::Continue) || matches!(s2, DFSStatus::Continue) {
644            DFSStatus::Continue
645        } else {
646            DFSStatus::Stop
647        }
648    }
649}
rapx/analysis/dataflow/types.rs

rapx/analysis/dataflow/
types.rs
