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

rapx/analysis/dataflow/
graph.rs
