| //! A pass that eliminates branches on uninhabited enum variants. |
| |
| use crate::transform::{MirPass, MirSource}; |
| use rustc::mir::{ |
| BasicBlock, BasicBlockData, Body, BodyCache, Local, Operand, Rvalue, StatementKind, |
| TerminatorKind, |
| }; |
| use rustc::ty::layout::{Abi, TyLayout, Variants}; |
| use rustc::ty::{Ty, TyCtxt}; |
| |
| pub struct UninhabitedEnumBranching; |
| |
| fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> { |
| if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator { |
| p.as_local() |
| } else { |
| None |
| } |
| } |
| |
| /// If the basic block terminates by switching on a discriminant, this returns the `Ty` the |
| /// discriminant is read from. Otherwise, returns None. |
| fn get_switched_on_type<'tcx>( |
| block_data: &BasicBlockData<'tcx>, |
| body: &Body<'tcx>, |
| ) -> Option<Ty<'tcx>> { |
| let terminator = block_data.terminator(); |
| |
| // Only bother checking blocks which terminate by switching on a local. |
| if let Some(local) = get_discriminant_local(&terminator.kind) { |
| let stmt_before_term = (block_data.statements.len() > 0) |
| .then_with(|| &block_data.statements[block_data.statements.len() - 1].kind); |
| |
| if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term |
| { |
| if l.as_local() == Some(local) { |
| if let Some(r_local) = place.as_local() { |
| let ty = body.local_decls[r_local].ty; |
| |
| if ty.is_enum() { |
| return Some(ty); |
| } |
| } |
| } |
| } |
| } |
| |
| None |
| } |
| |
| fn variant_discriminants<'tcx>( |
| layout: &TyLayout<'tcx>, |
| ty: Ty<'tcx>, |
| tcx: TyCtxt<'tcx>, |
| ) -> Vec<u128> { |
| match &layout.details.variants { |
| Variants::Single { index } => vec![index.as_u32() as u128], |
| Variants::Multiple { variants, .. } => variants |
| .iter_enumerated() |
| .filter_map(|(idx, layout)| { |
| (layout.abi != Abi::Uninhabited) |
| .then_with(|| ty.discriminant_for_variant(tcx, idx).unwrap().val) |
| }) |
| .collect(), |
| } |
| } |
| |
| impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching { |
| fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyCache<'tcx>) { |
| if source.promoted.is_some() { |
| return; |
| } |
| |
| trace!("UninhabitedEnumBranching starting for {:?}", source); |
| |
| let basic_block_count = body.basic_blocks().len(); |
| |
| for bb in 0..basic_block_count { |
| let bb = BasicBlock::from_usize(bb); |
| trace!("processing block {:?}", bb); |
| |
| let discriminant_ty = |
| if let Some(ty) = get_switched_on_type(&body.basic_blocks()[bb], body) { |
| ty |
| } else { |
| continue; |
| }; |
| |
| let layout = tcx.layout_of(tcx.param_env(source.def_id()).and(discriminant_ty)); |
| |
| let allowed_variants = if let Ok(layout) = layout { |
| variant_discriminants(&layout, discriminant_ty, tcx) |
| } else { |
| continue; |
| }; |
| |
| trace!("allowed_variants = {:?}", allowed_variants); |
| |
| if let TerminatorKind::SwitchInt { values, targets, .. } = |
| &mut body.basic_blocks_mut()[bb].terminator_mut().kind |
| { |
| let vals = &*values; |
| let zipped = vals.iter().zip(targets.into_iter()); |
| |
| let mut matched_values = Vec::with_capacity(allowed_variants.len()); |
| let mut matched_targets = Vec::with_capacity(allowed_variants.len() + 1); |
| |
| for (val, target) in zipped { |
| if allowed_variants.contains(val) { |
| matched_values.push(*val); |
| matched_targets.push(*target); |
| } else { |
| trace!("eliminating {:?} -> {:?}", val, target); |
| } |
| } |
| |
| // handle the "otherwise" branch |
| matched_targets.push(targets.pop().unwrap()); |
| |
| *values = matched_values.into(); |
| *targets = matched_targets; |
| } else { |
| unreachable!() |
| } |
| } |
| } |
| } |