compiler/rustc_middle/src/mir/patch.rs

   1 use rustc_index::vec::{Idx, IndexVec};
   2 use rustc_middle::mir::*;
   3 use rustc_middle::ty::Ty;
   4 use rustc_span::Span;
   5
   6 /// This struct represents a patch to MIR, which can add
   7 /// new statements and basic blocks and patch over block
   8 /// terminators.
   9 pub struct MirPatch<'tcx> {
  10     patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
  11     new_blocks: Vec<BasicBlockData<'tcx>>,
  12     new_statements: Vec<(Location, StatementKind<'tcx>)>,
  13     new_locals: Vec<LocalDecl<'tcx>>,
  14     resume_block: Option<BasicBlock>,
  15     body_span: Span,
  16     next_local: usize,
  17 }
  18
  19 impl<'tcx> MirPatch<'tcx> {
  20     pub fn new(body: &Body<'tcx>) -> Self {
  21         let mut result = MirPatch {
  22             patch_map: IndexVec::from_elem(None, body.basic_blocks()),
  23             new_blocks: vec![],
  24             new_statements: vec![],
  25             new_locals: vec![],
  26             next_local: body.local_decls.len(),
  27             resume_block: None,
  28             body_span: body.span,
  29         };
  30
  31         // Check if we already have a resume block
  32         for (bb, block) in body.basic_blocks().iter_enumerated() {
  33             if let TerminatorKind::Resume = block.terminator().kind && block.statements.is_empty() {
  34                 result.resume_block = Some(bb);
  35                 break;
  36             }
  37         }
  38
  39         result
  40     }
  41
  42     pub fn resume_block(&mut self) -> BasicBlock {
  43         if let Some(bb) = self.resume_block {
  44             return bb;
  45         }
  46
  47         let bb = self.new_block(BasicBlockData {
  48             statements: vec![],
  49             terminator: Some(Terminator {
  50                 source_info: SourceInfo::outermost(self.body_span),
  51                 kind: TerminatorKind::Resume,
  52             }),
  53             is_cleanup: true,
  54         });
  55         self.resume_block = Some(bb);
  56         bb
  57     }
  58
  59     pub fn is_patched(&self, bb: BasicBlock) -> bool {
  60         self.patch_map[bb].is_some()
  61     }
  62
  63     pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
  64         let offset = match bb.index().checked_sub(body.basic_blocks().len()) {
  65             Some(index) => self.new_blocks[index].statements.len(),
  66             None => body[bb].statements.len(),
  67         };
  68         Location { block: bb, statement_index: offset }
  69     }
  70
  71     pub fn new_local_with_info(
  72         &mut self,
  73         ty: Ty<'tcx>,
  74         span: Span,
  75         local_info: Option<Box<LocalInfo<'tcx>>>,
  76     ) -> Local {
  77         let index = self.next_local;
  78         self.next_local += 1;
  79         let mut new_decl = LocalDecl::new(ty, span);
  80         new_decl.local_info = local_info;
  81         self.new_locals.push(new_decl);
  82         Local::new(index as usize)
  83     }
  84
  85     pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
  86         self.new_local_with_info(ty, span, None)
  87     }
  88
  89     pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
  90         let index = self.next_local;
  91         self.next_local += 1;
  92         self.new_locals.push(LocalDecl::new(ty, span).internal());
  93         Local::new(index as usize)
  94     }
  95
  96     pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
  97         let block = BasicBlock::new(self.patch_map.len());
  98         debug!("MirPatch: new_block: {:?}: {:?}", block, data);
  99         self.new_blocks.push(data);
 100         self.patch_map.push(None);
 101         block
 102     }
 103
 104     pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
 105         assert!(self.patch_map[block].is_none());
 106         debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
 107         self.patch_map[block] = Some(new);
 108     }
 109
 110     pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
 111         debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
 112         self.new_statements.push((loc, stmt));
 113     }
 114
 115     pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
 116         self.add_statement(loc, StatementKind::Assign(Box::new((place, rv))));
 117     }
 118
 119     pub fn apply(self, body: &mut Body<'tcx>) {
 120         debug!(
 121             "MirPatch: {:?} new temps, starting from index {}: {:?}",
 122             self.new_locals.len(),
 123             body.local_decls.len(),
 124             self.new_locals
 125         );
 126         debug!(
 127             "MirPatch: {} new blocks, starting from index {}",
 128             self.new_blocks.len(),
 129             body.basic_blocks().len()
 130         );
 131         let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() {
 132             body.basic_blocks.as_mut_preserves_cfg()
 133         } else {
 134             body.basic_blocks.as_mut()
 135         };
 136         bbs.extend(self.new_blocks);
 137         body.local_decls.extend(self.new_locals);
 138         for (src, patch) in self.patch_map.into_iter_enumerated() {
 139             if let Some(patch) = patch {
 140                 debug!("MirPatch: patching block {:?}", src);
 141                 bbs[src].terminator_mut().kind = patch;
 142             }
 143         }
 144
 145         let mut new_statements = self.new_statements;
 146         new_statements.sort_by_key(|s| s.0);
 147
 148         let mut delta = 0;
 149         let mut last_bb = START_BLOCK;
 150         let mut stmts_and_targets: Vec<(Statement<'_>, BasicBlock)> = Vec::new();
 151         for (mut loc, stmt) in new_statements {
 152             if loc.block != last_bb {
 153                 delta = 0;
 154                 last_bb = loc.block;
 155             }
 156             debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
 157             loc.statement_index += delta;
 158             let source_info = Self::source_info_for_index(&body[loc.block], loc);
 159
 160             // For mir-opt `Derefer` to work in all cases we need to
 161             // get terminator's targets and apply the statement to all of them.
 162             if loc.statement_index > body[loc.block].statements.len() {
 163                 let term = body[loc.block].terminator();
 164                 for i in term.successors() {
 165                     stmts_and_targets.push((Statement { source_info, kind: stmt.clone() }, i));
 166                 }
 167                 delta += 1;
 168                 continue;
 169             }
 170
 171             body[loc.block]
 172                 .statements
 173                 .insert(loc.statement_index, Statement { source_info, kind: stmt });
 174             delta += 1;
 175         }
 176
 177         for (stmt, target) in stmts_and_targets.into_iter().rev() {
 178             body[target].statements.insert(0, stmt);
 179         }
 180     }
 181
 182     pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
 183         match data.statements.get(loc.statement_index) {
 184             Some(stmt) => stmt.source_info,
 185             None => data.terminator().source_info,
 186         }
 187     }
 188
 189     pub fn source_info_for_location(&self, body: &Body<'tcx>, loc: Location) -> SourceInfo {
 190         let data = match loc.block.index().checked_sub(body.basic_blocks().len()) {
 191             Some(new) => &self.new_blocks[new],
 192             None => &body[loc.block],
 193         };
 194         Self::source_info_for_index(data, loc)
 195     }
 196 }