|
| 1 | +//! This pass inserts the same validity checks into `MaybeUninit::{uninit,zeroed}().assert_init()` |
| 2 | +//! as in `mem::{uninitialized,zeroed}`. |
| 3 | +//! |
| 4 | +//! Note that this module uses `uninit` to mean `uninit` or `zeroed` unless `zeroed` is used explicitly. |
| 5 | +//! |
| 6 | +//! It does this by first finding a call to `MaybeUninit::uninit`, and then figuring out |
| 7 | +//! whether the successor basic block is a call to `MaybeUninit::assume_init` on the same local. |
| 8 | +
|
| 9 | +use rustc_const_eval::interpret; |
| 10 | +use rustc_hir::def_id::DefId; |
| 11 | +use rustc_middle::mir::patch::MirPatch; |
| 12 | +use rustc_middle::mir::{ |
| 13 | + BasicBlock, BasicBlockData, Body, Constant, ConstantKind, Operand, Place, SourceInfo, |
| 14 | + Terminator, TerminatorKind, |
| 15 | +}; |
| 16 | +use rustc_middle::ty::{self, List, SubstsRef, TyCtxt}; |
| 17 | +use rustc_span::{sym, Span}; |
| 18 | + |
| 19 | +use crate::MirPass; |
| 20 | + |
| 21 | +pub struct CheckMaybeUninit; |
| 22 | + |
| 23 | +impl<'tcx> MirPass<'tcx> for CheckMaybeUninit { |
| 24 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 25 | + let mut patch = MirPatch::new(body); |
| 26 | + |
| 27 | + for (mu_uninit_bb, _) in body.basic_blocks.iter_enumerated() { |
| 28 | + let terminator = body.basic_blocks[mu_uninit_bb].terminator(); |
| 29 | + |
| 30 | + let TerminatorKind::Call { |
| 31 | + func: mu_uninit_func, |
| 32 | + target: assume_init_bb, |
| 33 | + destination: uninit_place, |
| 34 | + .. |
| 35 | + } = &terminator.kind else { |
| 36 | + continue; |
| 37 | + }; |
| 38 | + |
| 39 | + let Some((mu_method_def_id, substs)) = mu_uninit_func.const_fn_def() else { |
| 40 | + continue; |
| 41 | + }; |
| 42 | + |
| 43 | + let Some(assume_init_bb) = assume_init_bb else { |
| 44 | + continue; |
| 45 | + }; |
| 46 | + |
| 47 | + let Some((assume_init_operand, assume_init_call_span)) = is_block_just_assume_init(tcx, &body.basic_blocks[*assume_init_bb]) else { |
| 48 | + continue; |
| 49 | + }; |
| 50 | + |
| 51 | + let Some(assume_init_place) = assume_init_operand.place() else { |
| 52 | + continue; |
| 53 | + }; |
| 54 | + |
| 55 | + if assume_init_place != *uninit_place { |
| 56 | + // The calls here are a little sketchy, but the place that is assumed to be init is not the place that was just crated |
| 57 | + // as uninit, so we conservatively bail out. |
| 58 | + continue; |
| 59 | + } |
| 60 | + |
| 61 | + // Select the right assertion intrinsic to call depending on which MaybeUninit method we called |
| 62 | + let Some(init_check_def_id) = get_init_check_def_id(tcx, mu_method_def_id) else { |
| 63 | + continue; |
| 64 | + }; |
| 65 | + |
| 66 | + let assert_valid_bb = make_assert_valid_bb( |
| 67 | + &mut patch, |
| 68 | + tcx, |
| 69 | + assume_init_call_span, |
| 70 | + init_check_def_id, |
| 71 | + *assume_init_bb, |
| 72 | + substs, |
| 73 | + ); |
| 74 | + |
| 75 | + let mut new_uninit_terminator = terminator.kind.clone(); |
| 76 | + match new_uninit_terminator { |
| 77 | + TerminatorKind::Call { ref mut target, .. } => { |
| 78 | + *target = Some(assert_valid_bb); |
| 79 | + } |
| 80 | + _ => unreachable!("terminator must be TerminatorKind::Call as checked above"), |
| 81 | + } |
| 82 | + |
| 83 | + patch.patch_terminator(mu_uninit_bb, new_uninit_terminator); |
| 84 | + } |
| 85 | + |
| 86 | + patch.apply(body); |
| 87 | + } |
| 88 | +} |
| 89 | + |
| 90 | +fn is_block_just_assume_init<'tcx, 'blk>( |
| 91 | + tcx: TyCtxt<'tcx>, |
| 92 | + block: &'blk BasicBlockData<'tcx>, |
| 93 | +) -> Option<(&'blk Operand<'tcx>, Span)> { |
| 94 | + if block.statements.is_empty() |
| 95 | + && let TerminatorKind::Call { |
| 96 | + func, |
| 97 | + args, |
| 98 | + fn_span, |
| 99 | + .. |
| 100 | + } = &block.terminator().kind |
| 101 | + && let Some((def_id, _)) = func.const_fn_def() |
| 102 | + && tcx.is_diagnostic_item(sym::assume_init, def_id) |
| 103 | + { |
| 104 | + args.get(0).map(|operand| (operand, *fn_span)) |
| 105 | + } else { |
| 106 | + None |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +fn get_init_check_def_id(tcx: TyCtxt<'_>, mu_method_def_id: DefId) -> Option<DefId> { |
| 111 | + if tcx.is_diagnostic_item(sym::maybe_uninit_uninit, mu_method_def_id) { |
| 112 | + tcx.lang_items().assert_uninit_valid() |
| 113 | + } else if tcx.is_diagnostic_item(sym::maybe_uninit_zeroed, mu_method_def_id) { |
| 114 | + tcx.lang_items().assert_zero_valid() |
| 115 | + } else { |
| 116 | + None |
| 117 | + } |
| 118 | +} |
| 119 | + |
| 120 | +fn make_assert_valid_bb<'tcx>( |
| 121 | + patch: &mut MirPatch<'tcx>, |
| 122 | + tcx: TyCtxt<'tcx>, |
| 123 | + fn_span: Span, |
| 124 | + init_check_def_id: DefId, |
| 125 | + target_bb: BasicBlock, |
| 126 | + substs: SubstsRef<'tcx>, |
| 127 | +) -> BasicBlock { |
| 128 | + let func = make_fn_operand_for_assert_valid(tcx, init_check_def_id, fn_span, substs); |
| 129 | + |
| 130 | + let local = patch.new_temp(tcx.types.unit, fn_span); |
| 131 | + |
| 132 | + let terminator = TerminatorKind::Call { |
| 133 | + func, |
| 134 | + args: vec![], |
| 135 | + destination: Place { local, projection: List::empty() }, |
| 136 | + target: Some(target_bb), |
| 137 | + cleanup: Some(patch.resume_block()), |
| 138 | + from_hir_call: true, |
| 139 | + fn_span, |
| 140 | + }; |
| 141 | + |
| 142 | + let terminator = Terminator { source_info: SourceInfo::outermost(fn_span), kind: terminator }; |
| 143 | + |
| 144 | + let bb_data = BasicBlockData::new(Some(terminator)); |
| 145 | + |
| 146 | + let block = patch.new_block(bb_data); |
| 147 | + block |
| 148 | +} |
| 149 | + |
| 150 | +fn make_fn_operand_for_assert_valid<'tcx>( |
| 151 | + tcx: TyCtxt<'tcx>, |
| 152 | + def_id: DefId, |
| 153 | + span: Span, |
| 154 | + substs: SubstsRef<'tcx>, |
| 155 | +) -> Operand<'tcx> { |
| 156 | + let fn_ty = ty::FnDef(def_id, substs); |
| 157 | + let fn_ty = tcx.mk_ty(fn_ty); |
| 158 | + |
| 159 | + Operand::Constant(Box::new(Constant { |
| 160 | + span, |
| 161 | + literal: ConstantKind::Val(interpret::ConstValue::ZeroSized, fn_ty), |
| 162 | + user_ty: None, |
| 163 | + })) |
| 164 | +} |
0 commit comments