rustc_mir_transform/
add_retag.rs

1//! This pass adds validation calls (AcquireValid, ReleaseValid) where appropriate.
2//! It has to be run really early, before transformations like inlining, because
3//! introducing these calls *adds* UB -- so, conceptually, this pass is actually part
4//! of MIR building, and only after this pass we think of the program has having the
5//! normal MIR semantics.
6
7use rustc_hir::LangItem;
8use rustc_middle::mir::*;
9use rustc_middle::ty::{self, Ty, TyCtxt};
10
11pub(super) struct AddRetag;
12
13/// Determine whether this type may contain a reference (or box), and thus needs retagging.
14/// We will only recurse `depth` times into Tuples/ADTs to bound the cost of this.
15fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> bool {
16    match ty.kind() {
17        // Primitive types that are not references
18        ty::Bool
19        | ty::Char
20        | ty::Float(_)
21        | ty::Int(_)
22        | ty::Uint(_)
23        | ty::RawPtr(..)
24        | ty::FnPtr(..)
25        | ty::Str
26        | ty::FnDef(..)
27        | ty::Never => false,
28        // References and Boxes (`noalias` sources)
29        ty::Ref(..) => true,
30        ty::Adt(..) if ty.is_box() => true,
31        ty::Adt(adt, _) if tcx.is_lang_item(adt.did(), LangItem::PtrUnique) => true,
32        // Compound types: recurse
33        ty::Array(ty, _) | ty::Slice(ty) => {
34            // This does not branch so we keep the depth the same.
35            may_contain_reference(*ty, depth, tcx)
36        }
37        ty::Tuple(tys) => {
38            depth == 0 || tys.iter().any(|ty| may_contain_reference(ty, depth - 1, tcx))
39        }
40        ty::Adt(adt, args) => {
41            depth == 0
42                || adt.variants().iter().any(|v| {
43                    v.fields.iter().any(|f| may_contain_reference(f.ty(tcx, args), depth - 1, tcx))
44                })
45        }
46        // Conservative fallback
47        _ => true,
48    }
49}
50
51impl<'tcx> crate::MirPass<'tcx> for AddRetag {
52    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
53        sess.opts.unstable_opts.mir_emit_retag
54    }
55
56    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
57        // We need an `AllCallEdges` pass before we can do any work.
58        super::add_call_guards::AllCallEdges.run_pass(tcx, body);
59
60        let basic_blocks = body.basic_blocks.as_mut();
61        let local_decls = &body.local_decls;
62        let needs_retag = |place: &Place<'tcx>| {
63            // We're not really interested in stores to "outside" locations, they are hard to keep
64            // track of anyway.
65            !place.is_indirect_first_projection()
66                && may_contain_reference(place.ty(&*local_decls, tcx).ty, /*depth*/ 3, tcx)
67                && !local_decls[place.local].is_deref_temp()
68        };
69
70        // PART 1
71        // Retag arguments at the beginning of the start block.
72        {
73            // Gather all arguments, skip return value.
74            let places = local_decls.iter_enumerated().skip(1).take(body.arg_count).filter_map(
75                |(local, decl)| {
76                    let place = Place::from(local);
77                    needs_retag(&place).then_some((place, decl.source_info))
78                },
79            );
80
81            // Emit their retags.
82            basic_blocks[START_BLOCK].statements.splice(
83                0..0,
84                places.map(|(place, source_info)| Statement {
85                    source_info,
86                    kind: StatementKind::Retag(RetagKind::FnEntry, Box::new(place)),
87                }),
88            );
89        }
90
91        // PART 2
92        // Retag return values of functions.
93        // We collect the return destinations because we cannot mutate while iterating.
94        let returns = basic_blocks
95            .iter_mut()
96            .filter_map(|block_data| {
97                match block_data.terminator().kind {
98                    TerminatorKind::Call { target: Some(target), destination, .. }
99                        if needs_retag(&destination) =>
100                    {
101                        // Remember the return destination for later
102                        Some((block_data.terminator().source_info, destination, target))
103                    }
104
105                    // `Drop` is also a call, but it doesn't return anything so we are good.
106                    TerminatorKind::Drop { .. } => None,
107                    // Not a block ending in a Call -> ignore.
108                    _ => None,
109                }
110            })
111            .collect::<Vec<_>>();
112        // Now we go over the returns we collected to retag the return values.
113        for (source_info, dest_place, dest_block) in returns {
114            basic_blocks[dest_block].statements.insert(
115                0,
116                Statement {
117                    source_info,
118                    kind: StatementKind::Retag(RetagKind::Default, Box::new(dest_place)),
119                },
120            );
121        }
122
123        // PART 3
124        // Add retag after assignments.
125        for block_data in basic_blocks {
126            // We want to insert statements as we iterate. To this end, we
127            // iterate backwards using indices.
128            for i in (0..block_data.statements.len()).rev() {
129                let (retag_kind, place) = match block_data.statements[i].kind {
130                    // Retag after assignments of reference type.
131                    StatementKind::Assign(box (ref place, ref rvalue)) => {
132                        let add_retag = match rvalue {
133                            // Ptr-creating operations already do their own internal retagging, no
134                            // need to also add a retag statement. *Except* if we are deref'ing a
135                            // Box, because those get desugared to directly working with the inner
136                            // raw pointer! That's relevant for `RawPtr` as Miri otherwise makes it
137                            // a NOP when the original pointer is already raw.
138                            Rvalue::RawPtr(_mutbl, place) => {
139                                // Using `is_box_global` here is a bit sketchy: if this code is
140                                // generic over the allocator, we'll not add a retag! This is a hack
141                                // to make Stacked Borrows compatible with custom allocator code.
142                                // It means the raw pointer inherits the tag of the box, which mostly works
143                                // but can sometimes lead to unexpected aliasing errors.
144                                // Long-term, we'll want to move to an aliasing model where "cast to
145                                // raw pointer" is a complete NOP, and then this will no longer be
146                                // an issue.
147                                if place.is_indirect_first_projection()
148                                    && body.local_decls[place.local].ty.is_box_global(tcx)
149                                {
150                                    Some(RetagKind::Raw)
151                                } else {
152                                    None
153                                }
154                            }
155                            Rvalue::Ref(..) => None,
156                            _ => {
157                                if needs_retag(place) {
158                                    Some(RetagKind::Default)
159                                } else {
160                                    None
161                                }
162                            }
163                        };
164                        if let Some(kind) = add_retag {
165                            (kind, *place)
166                        } else {
167                            continue;
168                        }
169                    }
170                    // Do nothing for the rest
171                    _ => continue,
172                };
173                // Insert a retag after the statement.
174                let source_info = block_data.statements[i].source_info;
175                block_data.statements.insert(
176                    i + 1,
177                    Statement {
178                        source_info,
179                        kind: StatementKind::Retag(retag_kind, Box::new(place)),
180                    },
181                );
182            }
183        }
184    }
185
186    fn is_required(&self) -> bool {
187        true
188    }
189}