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 /// 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.
  15 fn 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
  29         ty::Ref(..) => true,
  30         ty::Adt(..) if ty.is_box() => true,
  31         // Compound types: recurse
  32         ty::Array(ty, _) | ty::Slice(ty) => {
  33             // This does not branch so we keep the depth the same.
  34             may_contain_reference(*ty, depth, tcx)
  35         }
  36         ty::Tuple(tys) => {
  37             depth == 0 || tys.iter().any(|ty| may_contain_reference(ty, depth - 1, tcx))
  38         }
  39         ty::Adt(adt, subst) => {
  40             depth == 0
  41                 || adt.variants().iter().any(|v| {
  42                     v.fields.iter().any(|f| may_contain_reference(f.ty(tcx, subst), depth - 1, tcx))
  43                 })
  44         }
  45         // Conservative fallback
  46         _ => true,
  47     }
  48 }
  49
  50 impl<'tcx> MirPass<'tcx> for AddRetag {
  51     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  52         sess.opts.unstable_opts.mir_emit_retag
  53     }
  54
  55     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  56         // We need an `AllCallEdges` pass before we can do any work.
  57         super::add_call_guards::AllCallEdges.run_pass(tcx, body);
  58
  59         let basic_blocks = body.basic_blocks.as_mut();
  60         let local_decls = &body.local_decls;
  61         let needs_retag = |place: &Place<'tcx>| {
  62             !place.has_deref() // we're not eally interested in stores to "outside" locations, they are hard to keep track of anyway
  63                 && may_contain_reference(place.ty(&*local_decls, tcx).ty, /*depth*/ 3, tcx)
  64                 && !local_decls[place.local].is_deref_temp()
  65         };
  66
  67         // PART 1
  68         // Retag arguments at the beginning of the start block.
  69         {
  70             // Gather all arguments, skip return value.
  71             let places = local_decls.iter_enumerated().skip(1).take(body.arg_count).filter_map(
  72                 |(local, decl)| {
  73                     let place = Place::from(local);
  74                     needs_retag(&place).then_some((place, decl.source_info))
  75                 },
  76             );
  77
  78             // Emit their retags.
  79             basic_blocks[START_BLOCK].statements.splice(
  80                 0..0,
  81                 places.map(|(place, source_info)| Statement {
  82                     source_info,
  83                     kind: StatementKind::Retag(RetagKind::FnEntry, Box::new(place)),
  84                 }),
  85             );
  86         }
  87
  88         // PART 2
  89         // Retag return values of functions.
  90         // We collect the return destinations because we cannot mutate while iterating.
  91         let returns = basic_blocks
  92             .iter_mut()
  93             .filter_map(|block_data| {
  94                 match block_data.terminator().kind {
  95                     TerminatorKind::Call { target: Some(target), destination, .. }
  96                         if needs_retag(&destination) =>
  97                     {
  98                         // Remember the return destination for later
  99                         Some((block_data.terminator().source_info, destination, target))
 100                     }
 101
 102                     // `Drop` is also a call, but it doesn't return anything so we are good.
 103                     TerminatorKind::Drop { .. } | TerminatorKind::DropAndReplace { .. } => None,
 104                     // Not a block ending in a Call -> ignore.
 105                     _ => None,
 106                 }
 107             })
 108             .collect::<Vec<_>>();
 109         // Now we go over the returns we collected to retag the return values.
 110         for (source_info, dest_place, dest_block) in returns {
 111             basic_blocks[dest_block].statements.insert(
 112                 0,
 113                 Statement {
 114                     source_info,
 115                     kind: StatementKind::Retag(RetagKind::Default, Box::new(dest_place)),
 116                 },
 117             );
 118         }
 119
 120         // PART 3
 121         // Add retag after assignments where data "enters" this function: the RHS is behind a deref and the LHS is not.
 122         for block_data in basic_blocks {
 123             // We want to insert statements as we iterate. To this end, we
 124             // iterate backwards using indices.
 125             for i in (0..block_data.statements.len()).rev() {
 126                 let (retag_kind, place) = match block_data.statements[i].kind {
 127                     // Retag after assignments of reference type.
 128                     StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
 129                         let add_retag = match rvalue {
 130                             // Ptr-creating operations already do their own internal retagging, no
 131                             // need to also add a retag statement.
 132                             Rvalue::Ref(..) | Rvalue::AddressOf(..) => false,
 133                             _ => true,
 134                         };
 135                         if add_retag {
 136                             (RetagKind::Default, *place)
 137                         } else {
 138                             continue;
 139                         }
 140                     }
 141                     // Do nothing for the rest
 142                     _ => continue,
 143                 };
 144                 // Insert a retag after the statement.
 145                 let source_info = block_data.statements[i].source_info;
 146                 block_data.statements.insert(
 147                     i + 1,
 148                     Statement {
 149                         source_info,
 150                         kind: StatementKind::Retag(retag_kind, Box::new(place)),
 151                     },
 152                 );
 153             }
 154         }
 155     }
 156 }