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}