src/librustc_mir/util/patch.rs

   1 use rustc::mir::*;
   2 use rustc::ty::Ty;
   3 use rustc_index::vec::{Idx, IndexVec};
   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: BasicBlock,
  15     next_local: usize,
  16     make_nop: Vec<Location>,
  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: START_BLOCK,
  28             make_nop: vec![],
  29         };
  30
  31         // make sure the MIR we create has a resume block. It is
  32         // completely legal to convert jumps to the resume block
  33         // to jumps to None, but we occasionally have to add
  34         // instructions just before that.
  35
  36         let mut resume_block = None;
  37         let mut resume_stmt_block = None;
  38         for (bb, block) in body.basic_blocks().iter_enumerated() {
  39             if let TerminatorKind::Resume = block.terminator().kind {
  40                 if !block.statements.is_empty() {
  41                     assert!(resume_stmt_block.is_none());
  42                     resume_stmt_block = Some(bb);
  43                 } else {
  44                     resume_block = Some(bb);
  45                 }
  46                 break;
  47             }
  48         }
  49         let resume_block = resume_block.unwrap_or_else(|| {
  50             result.new_block(BasicBlockData {
  51                 statements: vec![],
  52                 terminator: Some(Terminator {
  53                     source_info: SourceInfo { span: body.span, scope: OUTERMOST_SOURCE_SCOPE },
  54                     kind: TerminatorKind::Resume,
  55                 }),
  56                 is_cleanup: true,
  57             })
  58         });
  59         result.resume_block = resume_block;
  60         if let Some(resume_stmt_block) = resume_stmt_block {
  61             result
  62                 .patch_terminator(resume_stmt_block, TerminatorKind::Goto { target: resume_block });
  63         }
  64         result
  65     }
  66
  67     pub fn resume_block(&self) -> BasicBlock {
  68         self.resume_block
  69     }
  70
  71     pub fn is_patched(&self, bb: BasicBlock) -> bool {
  72         self.patch_map[bb].is_some()
  73     }
  74
  75     pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
  76         let offset = match bb.index().checked_sub(body.basic_blocks().len()) {
  77             Some(index) => self.new_blocks[index].statements.len(),
  78             None => body[bb].statements.len(),
  79         };
  80         Location { block: bb, statement_index: offset }
  81     }
  82
  83     pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
  84         let index = self.next_local;
  85         self.next_local += 1;
  86         self.new_locals.push(LocalDecl::new_temp(ty, span));
  87         Local::new(index as usize)
  88     }
  89
  90     pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
  91         let index = self.next_local;
  92         self.next_local += 1;
  93         self.new_locals.push(LocalDecl::new_internal(ty, span));
  94         Local::new(index as usize)
  95     }
  96
  97     pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
  98         let block = BasicBlock::new(self.patch_map.len());
  99         debug!("MirPatch: new_block: {:?}: {:?}", block, data);
 100         self.new_blocks.push(data);
 101         self.patch_map.push(None);
 102         block
 103     }
 104
 105     pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
 106         assert!(self.patch_map[block].is_none());
 107         debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
 108         self.patch_map[block] = Some(new);
 109     }
 110
 111     pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
 112         debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
 113         self.new_statements.push((loc, stmt));
 114     }
 115
 116     pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
 117         self.add_statement(loc, StatementKind::Assign(box (place, rv)));
 118     }
 119
 120     pub fn make_nop(&mut self, loc: Location) {
 121         self.make_nop.push(loc);
 122     }
 123
 124     pub fn apply(self, body: &mut BodyAndCache<'tcx>) {
 125         debug!("MirPatch: make nops at: {:?}", self.make_nop);
 126         for loc in self.make_nop {
 127             body.make_statement_nop(loc);
 128         }
 129         debug!(
 130             "MirPatch: {:?} new temps, starting from index {}: {:?}",
 131             self.new_locals.len(),
 132             body.local_decls.len(),
 133             self.new_locals
 134         );
 135         debug!(
 136             "MirPatch: {} new blocks, starting from index {}",
 137             self.new_blocks.len(),
 138             body.basic_blocks().len()
 139         );
 140         body.basic_blocks_mut().extend(self.new_blocks);
 141         body.local_decls.extend(self.new_locals);
 142         for (src, patch) in self.patch_map.into_iter_enumerated() {
 143             if let Some(patch) = patch {
 144                 debug!("MirPatch: patching block {:?}", src);
 145                 body[src].terminator_mut().kind = patch;
 146             }
 147         }
 148
 149         let mut new_statements = self.new_statements;
 150         new_statements.sort_by_key(|s| s.0);
 151
 152         let mut delta = 0;
 153         let mut last_bb = START_BLOCK;
 154         for (mut loc, stmt) in new_statements {
 155             if loc.block != last_bb {
 156                 delta = 0;
 157                 last_bb = loc.block;
 158             }
 159             debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
 160             loc.statement_index += delta;
 161             let source_info = Self::source_info_for_index(&body[loc.block], loc);
 162             body[loc.block]
 163                 .statements
 164                 .insert(loc.statement_index, Statement { source_info, kind: stmt });
 165             delta += 1;
 166         }
 167     }
 168
 169     pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
 170         match data.statements.get(loc.statement_index) {
 171             Some(stmt) => stmt.source_info,
 172             None => data.terminator().source_info,
 173         }
 174     }
 175
 176     pub fn source_info_for_location(&self, body: &Body<'_>, loc: Location) -> SourceInfo {
 177         let data = match loc.block.index().checked_sub(body.basic_blocks().len()) {
 178             Some(new) => &self.new_blocks[new],
 179             None => &body[loc.block],
 180         };
 181         Self::source_info_for_index(data, loc)
 182     }
 183 }