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 _0 = Local::from_usize(0);
236        let deps = (0..self.argc + 1)
237            .map(|i| {
238                let _i = Local::from_usize(i);
239                self.is_connected(_i, _0)
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(start, direction, &mut node_op, edge_validator, false, &mut seen);
323        result
324    }
325
326    pub fn find_all_nodes<P, E>(
327        &self,
328        start: Local,
329        direction: Direction,
330        node_predicate: &mut P,
331        edge_validator: &mut E,
332    ) -> Vec<Local>
333    where
334        P: FnMut(&DataflowGraph, Local) -> bool,
335        E: FnMut(&DataflowGraph, EdgeIdx) -> DFSStatus,
336    {
337        let mut results = Vec::new();
338        let mut node_op = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
339            if node_predicate(graph, idx) {
340                results.push(idx);
341            }
342            DFSStatus::Continue
343        };
344        let mut seen = HashSet::new();
345        self.dfs(start, direction, &mut node_op, edge_validator, true, &mut seen);
346        results
347    }
348
349    pub fn equivalent_edge_validator(graph: &DataflowGraph, idx: EdgeIdx) -> DFSStatus {
350        match graph.edges[idx].op {
351            EdgeOp::Copy | EdgeOp::Move | EdgeOp::Mut | EdgeOp::Immut | EdgeOp::Deref => {
352                DFSStatus::Continue
353            }
354            EdgeOp::Nop
355            | EdgeOp::Const
356            | EdgeOp::Downcast(_)
357            | EdgeOp::Field(_)
358            | EdgeOp::Index
359            | EdgeOp::ConstIndex
360            | EdgeOp::SubSlice => DFSStatus::Stop,
361        }
362    }
363
364    pub fn always_true_edge_validator(_: &DataflowGraph, _: EdgeIdx) -> DFSStatus {
365        DFSStatus::Continue
366    }
367
368    pub fn collect_equivalent_locals(&self, local: Local, strict: bool) -> HashSet<Local> {
369        let mut set = HashSet::new();
370        let root = Cell::new(local);
371        let reduce_func = if strict {
372            DFSStatus::and
373        } else {
374            DFSStatus::or
375        };
376        let mut find_root_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
377            let node = &graph.nodes[idx];
378            node.ops
379                .iter()
380                .map(|op| match op {
381                    NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
382                        root.set(idx);
383                        DFSStatus::Continue
384                    }
385                    NodeOp::Call(_) => {
386                        root.set(idx);
387                        DFSStatus::Stop
388                    }
389                    _ => DFSStatus::Stop,
390                })
391                .reduce(reduce_func)
392                .unwrap()
393        };
394        let mut find_equivalent_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
395            let node = &graph.nodes[idx];
396            if set.contains(&idx) {
397                return DFSStatus::Stop;
398            }
399            node.ops
400                .iter()
401                .map(|op| match op {
402                    NodeOp::Nop | NodeOp::Use | NodeOp::Ref => {
403                        set.insert(idx);
404                        DFSStatus::Continue
405                    }
406                    NodeOp::Call(_) => {
407                        if idx == root.get() {
408                            set.insert(idx);
409                            DFSStatus::Continue
410                        } else {
411                            DFSStatus::Stop
412                        }
413                    }
414                    _ => DFSStatus::Stop,
415                })
416                .reduce(reduce_func)
417                .unwrap()
418        };
419        let mut seen = HashSet::new();
420        self.dfs(
421            local,
422            Direction::Upside,
423            &mut find_root_operator,
424            &mut Self::equivalent_edge_validator,
425            true,
426            &mut seen,
427        );
428        seen.clear();
429        self.dfs(
430            root.get(),
431            Direction::Downside,
432            &mut find_equivalent_operator,
433            &mut Self::equivalent_edge_validator,
434            true,
435            &mut seen,
436        );
437        set
438    }
439
440    fn collect_by_direction(
441        &self,
442        local: Local,
443        self_included: bool,
444        direction: Direction,
445    ) -> HashSet<Local> {
446        let mut ret = HashSet::new();
447        let mut node_operator = |_: &DataflowGraph, idx: Local| -> DFSStatus {
448            ret.insert(idx);
449            DFSStatus::Continue
450        };
451        let mut seen = HashSet::new();
452        self.dfs(
453            local,
454            direction,
455            &mut node_operator,
456            &mut DataflowGraph::always_true_edge_validator,
457            true,
458            &mut seen,
459        );
460        if !self_included {
461            ret.remove(&local);
462        }
463        ret
464    }
465
466    pub fn collect_ancestor_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
467        self.collect_by_direction(local, self_included, Direction::Upside)
468    }
469
470    pub fn collect_descending_locals(&self, local: Local, self_included: bool) -> HashSet<Local> {
471        self.collect_by_direction(local, self_included, Direction::Downside)
472    }
473
474    pub fn get_field_sequence(&self, local: Local) -> Option<(Local, Vec<usize>)> {
475        let mut fields = vec![];
476        let var = Cell::new(local);
477        let mut node_operator = |graph: &DataflowGraph, idx: Local| -> DFSStatus {
478            if graph.is_marker(idx) {
479                DFSStatus::Continue
480            } else {
481                var.set(idx);
482                DFSStatus::Stop
483            }
484        };
485        let mut edge_validator = |graph: &DataflowGraph, idx: EdgeIdx| -> DFSStatus {
486            if let EdgeOp::Field(field) = graph.edges[idx].op {
487                fields.insert(0, field);
488                DFSStatus::Continue
489            } else {
490                DFSStatus::Stop
491            }
492        };
493        let mut seen = HashSet::new();
494        self.dfs(
495            local,
496            Direction::Upside,
497            &mut node_operator,
498            &mut edge_validator,
499            false,
500            &mut seen,
501        );
502        if fields.is_empty() {
503            None
504        } else {
505            Some((var.get(), fields))
506        }
507    }
508
509    /// Follow Copy/Move edges upward to find the root real local behind any
510    /// copy chains (no projections). Stops at 16 hops to bound cycles.
511    pub fn trace_origin(&self, local: Local) -> Local {
512        let mut current = local;
513        let mut seen = HashSet::new();
514        for _ in 0..16 {
515            if !seen.insert(current) {
516                break;
517            }
518            let next = self.nodes[current]
519                .in_edges
520                .iter()
521                .find_map(|&ei| {
522                    let e = &self.edges[ei];
523                    if matches!(e.op, EdgeOp::Copy | EdgeOp::Move)
524                        && !self.is_marker(e.src)
525                    {
526                        Some(e.src)
527                    } else {
528                        None
529                    }
530                });
531            match next {
532                Some(src) if src != current => current = src,
533                _ => break,
534            }
535        }
536        current
537    }
538
539    /// Return `true` when `local` originates from a tuple field destructuring
540    /// (e.g. `(tuple.0, tuple.1)` after a call returning a tuple).
541    /// Follows Copy/Move chains upward through marker nodes and checks for
542    /// any `Field` edge along the projection chain.
543    pub fn is_from_tuple_field(&self, local: Local) -> bool {
544        let mut current = local;
545        let mut seen = HashSet::new();
546        for _ in 0..8 {
547            if !seen.insert(current) {
548                return false;
549            }
550            let mut next_local = None;
551            for &ei in &self.nodes[current].in_edges {
552                let e = &self.edges[ei];
553                if !matches!(e.op, EdgeOp::Copy | EdgeOp::Move) {
554                    continue;
555                }
556                if self.is_marker(e.src) {
557                    if self.marker_chain_has_field(e.src) {
558                        return true;
559                    }
560                    if let Some(real) = self.marker_to_real(e.src) {
561                        next_local = Some(real);
562                        break;
563                    }
564                } else {
565                    next_local = Some(e.src);
566                    break;
567                }
568            }
569            match next_local {
570                Some(src) if src != current => current = src,
571                _ => return false,
572            }
573        }
574        false
575    }
576
577    /// Walk up a projection-marker chain to find the underlying real local.
578    fn marker_to_real(&self, marker: Local) -> Option<Local> {
579        let mut current = marker;
580        for _ in 0..8 {
581            if !self.is_marker(current) {
582                return Some(current);
583            }
584            current = self.nodes[current]
585                .in_edges
586                .first()
587                .map(|&ei| self.edges[ei].src)?;
588        }
589        None
590    }
591
592    /// Check whether a projection-marker chain contains a `Field` edge.
593    fn marker_chain_has_field(&self, marker: Local) -> bool {
594        let mut current = marker;
595        for _ in 0..8 {
596            if !self.is_marker(current) {
597                return false;
598            }
599            let ei = match self.nodes[current].in_edges.first() {
600                Some(&ei) => ei,
601                None => return false,
602            };
603            if matches!(self.edges[ei].op, EdgeOp::Field(_)) {
604                return true;
605            }
606            current = self.edges[ei].src;
607        }
608        false
609    }
610}
611
612#[derive(Clone, Copy)]
613pub enum Direction {
614    Upside,
615    Downside,
616    Both,
617}
618
619pub enum DFSStatus {
620    Continue,
621    Stop,
622}
623
624impl DFSStatus {
625    pub fn and(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
626        if matches!(s1, DFSStatus::Stop) || matches!(s2, DFSStatus::Stop) {
627            DFSStatus::Stop
628        } else {
629            DFSStatus::Continue
630        }
631    }
632
633    pub fn or(s1: DFSStatus, s2: DFSStatus) -> DFSStatus {
634        if matches!(s1, DFSStatus::Continue) || matches!(s2, DFSStatus::Continue) {
635            DFSStatus::Continue
636        } else {
637            DFSStatus::Stop
638        }
639    }
640}
rapx/analysis/dataflow/types.rs

rapx/analysis/dataflow/
types.rs
