rapx/analysis/dataflow/
graph.rs

1use rustc_hir::def_id::DefId;
2use rustc_middle::{
3    mir::{
4        AggregateKind, BorrowKind, Const, Local, Operand, Place, PlaceElem, Rvalue, Statement,
5        StatementKind, Terminator, TerminatorKind,
6    },
7    ty::{TyCtxt, TyKind},
8};
9use rustc_span::Span;
10
11use super::types::*;
12
13/// Build a `DataflowGraph` for a single function identified by `def_id`.
14pub fn build_dataflow_graph(tcx: TyCtxt<'_>, def_id: DefId) -> DataflowGraph {
15    let body = tcx.optimized_mir(def_id);
16    build_dataflow_graph_from_body(def_id, body)
17}
18
19/// Build a `DataflowGraph` from a pre-existing MIR body (e.g. after SSA transformation).
20pub fn build_dataflow_graph_from_body(
21    def_id: DefId,
22    body: &rustc_middle::mir::Body<'_>,
23) -> DataflowGraph {
24    let mut graph = DataflowGraph::new(def_id, body.span, body.arg_count, body.local_decls.len());
25    for (block_idx, bb) in body.basic_blocks.iter().enumerate() {
26        for (stmt_idx, stmt) in bb.statements.iter().enumerate() {
27            graph.add_statm_to_graph(&stmt, block_idx, stmt_idx);
28        }
29        if let Some(terminator) = &bb.terminator {
30            let stmt_idx = bb.statements.len();
31            graph.add_terminator_to_graph(&terminator, block_idx, stmt_idx);
32        }
33    }
34    graph
35}
36
37impl DataflowGraph {
38    pub fn add_operand(&mut self, operand: &Operand, dst: Local, block: usize, stmt_idx: usize) {
39        match operand {
40            Operand::Copy(place) => {
41                let src = self.parse_place(place, block, stmt_idx);
42                self.add_node_edge(src, dst, EdgeOp::Copy, block, stmt_idx);
43            }
44            Operand::Move(place) => {
45                let src = self.parse_place(place, block, stmt_idx);
46                self.add_node_edge(src, dst, EdgeOp::Move, block, stmt_idx);
47            }
48            Operand::Constant(boxed_const_op) => {
49                let src_desc = boxed_const_op.const_.to_string();
50                let src_ty = match boxed_const_op.const_ {
51                    Const::Val(_, ty) => ty.to_string(),
52                    Const::Unevaluated(_, ty) => ty.to_string(),
53                    Const::Ty(ty, _) => ty.to_string(),
54                };
55                self.add_const_edge(src_desc, src_ty, dst, EdgeOp::Const, block, stmt_idx);
56            }
57            #[cfg(rapx_ge_99)]
58            Operand::RuntimeChecks(_) => {}
59        }
60    }
61
62    pub fn parse_place(&mut self, place: &Place, block: usize, stmt_idx: usize) -> Local {
63        fn parse_one_step(
64            graph: &mut DataflowGraph,
65            src: Local,
66            place_elem: PlaceElem,
67            block: usize,
68            stmt_idx: usize,
69        ) -> Local {
70            let dst = graph.nodes.push(DataflowNode::new());
71            match place_elem {
72                PlaceElem::Deref => {
73                    graph.add_node_edge(src, dst, EdgeOp::Deref, block, stmt_idx);
74                }
75                PlaceElem::Field(field_idx, _) => {
76                    graph.add_node_edge(src, dst, EdgeOp::Field(field_idx.as_usize()), block, stmt_idx);
77                }
78                PlaceElem::Downcast(symbol, _) => {
79                    graph.add_node_edge(src, dst, EdgeOp::Downcast(symbol.unwrap().to_string()), block, stmt_idx);
80                }
81                PlaceElem::Index(idx) => {
82                    graph.add_node_edge(src, dst, EdgeOp::Index, block, stmt_idx);
83                    graph.add_node_edge(idx, dst, EdgeOp::Nop, block, stmt_idx);
84                }
85                PlaceElem::ConstantIndex { .. } => {
86                    graph.add_node_edge(src, dst, EdgeOp::ConstIndex, block, stmt_idx);
87                }
88                PlaceElem::Subslice { .. } => {
89                    graph.add_node_edge(src, dst, EdgeOp::SubSlice, block, stmt_idx);
90                }
91                _ => {
92                    rap_debug!("{:?}", place_elem);
93                    todo!()
94                }
95            }
96            dst
97        }
98        let mut ret = place.local;
99        for place_elem in place.projection {
100            ret = parse_one_step(self, ret, place_elem, block, stmt_idx);
101        }
102        ret
103    }
104
105    pub fn add_statm_to_graph(&mut self, statement: &Statement, block: usize, stmt_idx: usize) {
106        if let StatementKind::Assign(boxed_statm) = &statement.kind {
107            let place = boxed_statm.0;
108            let dst = self.parse_place(&place, block, stmt_idx);
109            self.nodes[dst].span = statement.source_info.span;
110            let rvalue = &boxed_statm.1;
111            let seq = self.nodes[dst].seq;
112            if seq == self.nodes[dst].ops.len() {
113                self.nodes[dst].ops.push(NodeOp::Nop);
114            }
115            match rvalue {
116                Rvalue::Use(op, ..) => {
117                    self.add_operand(op, dst, block, stmt_idx);
118                    self.nodes[dst].ops[seq] = NodeOp::Use;
119                }
120                Rvalue::Repeat(op, _) => {
121                    self.add_operand(op, dst, block, stmt_idx);
122                    self.nodes[dst].ops[seq] = NodeOp::Repeat;
123                }
124                Rvalue::Ref(_, borrow_kind, place) => {
125                    let op = match borrow_kind {
126                        BorrowKind::Shared => EdgeOp::Immut,
127                        BorrowKind::Mut { .. } => EdgeOp::Mut,
128                        BorrowKind::Fake(_) => EdgeOp::Nop,
129                    };
130                    let src = self.parse_place(place, block, stmt_idx);
131                    self.add_node_edge(src, dst, op, block, stmt_idx);
132                    self.nodes[dst].ops[seq] = NodeOp::Ref;
133                }
134                Rvalue::Cast(_cast_kind, operand, _) => {
135                    self.add_operand(operand, dst, block, stmt_idx);
136                    self.nodes[dst].ops[seq] = NodeOp::Cast;
137                }
138                Rvalue::BinaryOp(_, operands) => {
139                    self.add_operand(&operands.0, dst, block, stmt_idx);
140                    self.add_operand(&operands.1, dst, block, stmt_idx);
141                    self.nodes[dst].ops[seq] = NodeOp::CheckedBinaryOp;
142                }
143                Rvalue::Aggregate(boxed_kind, operands) => {
144                    for operand in operands.iter() {
145                        self.add_operand(operand, dst, block, stmt_idx);
146                    }
147                    match **boxed_kind {
148                        AggregateKind::Array(_) => {
149                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::Array)
150                        }
151                        AggregateKind::Tuple => {
152                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::Tuple)
153                        }
154                        AggregateKind::Adt(def_id, ..) => {
155                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::Adt(def_id))
156                        }
157                        AggregateKind::Closure(def_id, ..) => {
158                            self.closures.insert(def_id);
159                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::Closure(def_id))
160                        }
161                        AggregateKind::Coroutine(def_id, ..) => {
162                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::Coroutine(def_id))
163                        }
164                        AggregateKind::RawPtr(_, _mutability) => {
165                            self.nodes[dst].ops[seq] = NodeOp::Aggregate(AggKind::RawPtr)
166                        }
167                        _ => {
168                            rap_debug!("{:?}", boxed_kind);
169                            todo!()
170                        }
171                    }
172                }
173                Rvalue::UnaryOp(_, operand) => {
174                    self.add_operand(operand, dst, block, stmt_idx);
175                    self.nodes[dst].ops[seq] = NodeOp::UnaryOp;
176                }
177                #[cfg(not(rapx_ge_99))]
178                Rvalue::NullaryOp(_) => {
179                    self.nodes[dst].ops[seq] = NodeOp::NullaryOp;
180                }
181                Rvalue::ThreadLocalRef(_) => {}
182                Rvalue::Discriminant(place) => {
183                    let src = self.parse_place(place, block, stmt_idx);
184                    self.add_node_edge(src, dst, EdgeOp::Nop, block, stmt_idx);
185                    self.nodes[dst].ops[seq] = NodeOp::Discriminant;
186                }
187                #[cfg(not(rapx_ge_99))]
188                Rvalue::ShallowInitBox(operand, _) => {
189                    self.add_operand(operand, dst, block, stmt_idx);
190                    self.nodes[dst].ops[seq] = NodeOp::ShallowInitBox;
191                }
192                Rvalue::CopyForDeref(place) => {
193                    let src = self.parse_place(place, block, stmt_idx);
194                    self.add_node_edge(src, dst, EdgeOp::Nop, block, stmt_idx);
195                    self.nodes[dst].ops[seq] = NodeOp::CopyForDeref;
196                }
197                Rvalue::RawPtr(_, place) => {
198                    let src = self.parse_place(place, block, stmt_idx);
199                    self.add_node_edge(src, dst, EdgeOp::Nop, block, stmt_idx);
200                    self.nodes[dst].ops[seq] = NodeOp::RawPtr;
201                }
202                _ => todo!(),
203            };
204            self.nodes[dst].seq = seq + 1;
205        }
206    }
207
208    pub fn add_terminator_to_graph(
209        &mut self,
210        terminator: &Terminator,
211        block: usize,
212        stmt_idx: usize,
213    ) {
214        if let TerminatorKind::Call {
215            func,
216            args,
217            destination,
218            ..
219        } = &terminator.kind
220        {
221            let dst = destination.local;
222            let seq = self.nodes[dst].seq;
223            if seq == self.nodes[dst].ops.len() {
224                self.nodes[dst].ops.push(NodeOp::Nop);
225            }
226            match func {
227                Operand::Constant(boxed_cnst) => {
228                    if let Const::Val(_, ty) = boxed_cnst.const_ {
229                        if let TyKind::FnDef(def_id, _) = ty.kind() {
230                            for op in args.iter() {
231                                self.add_operand(&op.node, dst, block, stmt_idx);
232                            }
233                            self.nodes[dst].ops[seq] = NodeOp::Call(*def_id);
234                        }
235                    }
236                }
237                Operand::Move(_) => {
238                    self.add_operand(func, dst, block, stmt_idx);
239                    for op in args.iter() {
240                        self.add_operand(&op.node, dst, block, stmt_idx);
241                    }
242                    self.nodes[dst].ops[seq] = NodeOp::CallOperand;
243                }
244                _ => {
245                    rap_debug!("{:?}", func);
246                    todo!();
247                }
248            }
249            self.nodes[dst].span = terminator.source_info.span;
250            self.nodes[dst].seq = seq + 1;
251        }
252    }
253
254    pub fn query_node_by_span(&self, span: Span, strict: bool) -> Option<(Local, &DataflowNode)> {
255        for (node_idx, node) in self.nodes.iter_enumerated() {
256            if strict {
257                if node.span == span {
258                    return Some((node_idx, node));
259                }
260            } else {
261                if !crate::utils::span::relative_pos_range(node.span, span).eq(0..0)
262                    && (node.span.lo() == span.lo() || node.span.hi() == span.hi())
263                {
264                    return Some((node_idx, node));
265                }
266            }
267        }
268        None
269    }
270}
rapx/analysis/dataflow/graph.rs

rapx/analysis/dataflow/
graph.rs
