compiler/rustc_mir_transform/src/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
   7 use crate::MirPass;
   8 use rustc_middle::mir::*;
   9 use rustc_middle::ty::{self, Ty, TyCtxt};
  10
  11 pub struct AddRetag;
  12
  13 /// Determines whether this place is "stable": Whether, if we evaluate it again
  14 /// after the assignment, we can be sure to obtain the same place value.
  15 /// (Concurrent accesses by other threads are no problem as these are anyway non-atomic
  16 /// copies.  Data races are UB.)
  17 fn is_stable(place: PlaceRef<'_>) -> bool {
  18     // Which place this evaluates to can change with any memory write,
  19     // so cannot assume deref to be stable.
  20     !place.has_deref()
  21 }
  22
  23 /// Determine whether this type may contain a reference (or box), and thus needs retagging.
  24 /// We will only recurse `depth` times into Tuples/ADTs to bound the cost of this.
  25 fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> bool {
  26     match ty.kind() {
  27         // Primitive types that are not references
  28         ty::Bool
  29         | ty::Char
  30         | ty::Float(_)
  31         | ty::Int(_)
  32         | ty::Uint(_)
  33         | ty::RawPtr(..)
  34         | ty::FnPtr(..)
  35         | ty::Str
  36         | ty::FnDef(..)
  37         | ty::Never => false,
  38         // References
  39         ty::Ref(..) => true,
  40         ty::Adt(..) if ty.is_box() => true,
  41         // Compound types: recurse
  42         ty::Array(ty, _) | ty::Slice(ty) => {
  43             // This does not branch so we keep the depth the same.
  44             may_contain_reference(*ty, depth, tcx)
  45         }
  46         ty::Tuple(tys) => {
  47             depth == 0 || tys.iter().any(|ty| may_contain_reference(ty, depth - 1, tcx))
  48         }
  49         ty::Adt(adt, subst) => {
  50             depth == 0
  51                 || adt.variants().iter().any(|v| {
  52                     v.fields.iter().any(|f| may_contain_reference(f.ty(tcx, subst), depth - 1, tcx))
  53                 })
  54         }
  55         // Conservative fallback
  56         _ => true,
  57     }
  58 }
  59
  60 impl<'tcx> MirPass<'tcx> for AddRetag {
  61     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  62         sess.opts.unstable_opts.mir_emit_retag
  63     }
  64
  65     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  66         // We need an `AllCallEdges` pass before we can do any work.
  67         super::add_call_guards::AllCallEdges.run_pass(tcx, body);
  68
  69         let basic_blocks = body.basic_blocks.as_mut();
  70         let local_decls = &body.local_decls;
  71         let needs_retag = |place: &Place<'tcx>| {
  72             // FIXME: Instead of giving up for unstable places, we should introduce
  73             // a temporary and retag on that.
  74             is_stable(place.as_ref())
  75                 && may_contain_reference(place.ty(&*local_decls, tcx).ty, /*depth*/ 3, tcx)
  76                 && !local_decls[place.local].is_deref_temp()
  77         };
  78         let place_base_raw = |place: &Place<'tcx>| {
  79             // If this is a `Deref`, get the type of what we are deref'ing.
  80             if place.has_deref() {
  81                 let ty = &local_decls[place.local].ty;
  82                 ty.is_unsafe_ptr()
  83             } else {
  84                 // Not a deref, and thus not raw.
  85                 false
  86             }
  87         };
  88
  89         // PART 1
  90         // Retag arguments at the beginning of the start block.
  91         {
  92             // Gather all arguments, skip return value.
  93             let places = local_decls.iter_enumerated().skip(1).take(body.arg_count).filter_map(
  94                 |(local, decl)| {
  95                     let place = Place::from(local);
  96                     needs_retag(&place).then_some((place, decl.source_info))
  97                 },
  98             );
  99
 100             // Emit their retags.
 101             basic_blocks[START_BLOCK].statements.splice(
 102                 0..0,
 103                 places.map(|(place, source_info)| Statement {
 104                     source_info,
 105                     kind: StatementKind::Retag(RetagKind::FnEntry, Box::new(place)),
 106                 }),
 107             );
 108         }
 109
 110         // PART 2
 111         // Retag return values of functions.  Also escape-to-raw the argument of `drop`.
 112         // We collect the return destinations because we cannot mutate while iterating.
 113         let returns = basic_blocks
 114             .iter_mut()
 115             .filter_map(|block_data| {
 116                 match block_data.terminator().kind {
 117                     TerminatorKind::Call { target: Some(target), destination, .. }
 118                         if needs_retag(&destination) =>
 119                     {
 120                         // Remember the return destination for later
 121                         Some((block_data.terminator().source_info, destination, target))
 122                     }
 123
 124                     // `Drop` is also a call, but it doesn't return anything so we are good.
 125                     TerminatorKind::Drop { .. } | TerminatorKind::DropAndReplace { .. } => None,
 126                     // Not a block ending in a Call -> ignore.
 127                     _ => None,
 128                 }
 129             })
 130             .collect::<Vec<_>>();
 131         // Now we go over the returns we collected to retag the return values.
 132         for (source_info, dest_place, dest_block) in returns {
 133             basic_blocks[dest_block].statements.insert(
 134                 0,
 135                 Statement {
 136                     source_info,
 137                     kind: StatementKind::Retag(RetagKind::Default, Box::new(dest_place)),
 138                 },
 139             );
 140         }
 141
 142         // PART 3
 143         // Add retag after assignment.
 144         for block_data in basic_blocks {
 145             // We want to insert statements as we iterate.  To this end, we
 146             // iterate backwards using indices.
 147             for i in (0..block_data.statements.len()).rev() {
 148                 let (retag_kind, place) = match block_data.statements[i].kind {
 149                     // Retag-as-raw after escaping to a raw pointer, if the referent
 150                     // is not already a raw pointer.
 151                     StatementKind::Assign(box (lplace, Rvalue::AddressOf(_, ref rplace)))
 152                         if !place_base_raw(rplace) =>
 153                     {
 154                         (RetagKind::Raw, lplace)
 155                     }
 156                     // Retag after assignments of reference type.
 157                     StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
 158                         let kind = match rvalue {
 159                             Rvalue::Ref(_, borrow_kind, _)
 160                                 if borrow_kind.allows_two_phase_borrow() =>
 161                             {
 162                                 RetagKind::TwoPhase
 163                             }
 164                             _ => RetagKind::Default,
 165                         };
 166                         (kind, *place)
 167                     }
 168                     // Do nothing for the rest
 169                     _ => continue,
 170                 };
 171                 // Insert a retag after the statement.
 172                 let source_info = block_data.statements[i].source_info;
 173                 block_data.statements.insert(
 174                     i + 1,
 175                     Statement {
 176                         source_info,
 177                         kind: StatementKind::Retag(retag_kind, Box::new(place)),
 178                     },
 179                 );
 180             }
 181         }
 182     }
 183 }