src/librustc_mir/util/patch.rs

   1 use rustc::ty::Ty;
   2 use rustc::mir::*;
   3 use rustc_data_structures::indexed_vec::{IndexVec, Idx};
   4 use syntax_pos::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(mir: &Mir<'tcx>) -> Self {
  21         let mut result = MirPatch {
  22             patch_map: IndexVec::from_elem(None, mir.basic_blocks()),
  23             new_blocks: vec![],
  24             new_statements: vec![],
  25             new_locals: vec![],
  26             next_local: mir.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 mir.basic_blocks().iter_enumerated() {
  39             if let TerminatorKind::Resume = block.terminator().kind {
  40                 if block.statements.len() > 0 {
  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 {
  54                         span: mir.span,
  55                         scope: OUTERMOST_SOURCE_SCOPE
  56                     },
  57                     kind: TerminatorKind::Resume
  58                 }),
  59                 is_cleanup: true
  60             })});
  61         result.resume_block = resume_block;
  62         if let Some(resume_stmt_block) = resume_stmt_block {
  63             result.patch_terminator(resume_stmt_block, TerminatorKind::Goto {
  64                 target: resume_block
  65             });
  66         }
  67         result
  68     }
  69
  70     pub fn resume_block(&self) -> BasicBlock {
  71         self.resume_block
  72     }
  73
  74     pub fn is_patched(&self, bb: BasicBlock) -> bool {
  75         self.patch_map[bb].is_some()
  76     }
  77
  78     pub fn terminator_loc(&self, mir: &Mir<'tcx>, bb: BasicBlock) -> Location {
  79         let offset = match bb.index().checked_sub(mir.basic_blocks().len()) {
  80             Some(index) => self.new_blocks[index].statements.len(),
  81             None => mir[bb].statements.len()
  82         };
  83         Location {
  84             block: bb,
  85             statement_index: offset
  86         }
  87     }
  88
  89     pub fn new_temp(&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_temp(ty, span));
  93         Local::new(index as usize)
  94     }
  95
  96     pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
  97         let index = self.next_local;
  98         self.next_local += 1;
  99         self.new_locals.push(LocalDecl::new_internal(ty, span));
 100         Local::new(index as usize)
 101     }
 102
 103     pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
 104         let block = BasicBlock::new(self.patch_map.len());
 105         debug!("MirPatch: new_block: {:?}: {:?}", block, data);
 106         self.new_blocks.push(data);
 107         self.patch_map.push(None);
 108         block
 109     }
 110
 111     pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
 112         assert!(self.patch_map[block].is_none());
 113         debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
 114         self.patch_map[block] = Some(new);
 115     }
 116
 117     pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
 118         debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
 119         self.new_statements.push((loc, stmt));
 120     }
 121
 122     pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
 123         self.add_statement(loc, StatementKind::Assign(place, box rv));
 124     }
 125
 126     pub fn make_nop(&mut self, loc: Location) {
 127         self.make_nop.push(loc);
 128     }
 129
 130     pub fn apply(self, mir: &mut Mir<'tcx>) {
 131         debug!("MirPatch: make nops at: {:?}", self.make_nop);
 132         for loc in self.make_nop {
 133             mir.make_statement_nop(loc);
 134         }
 135         debug!("MirPatch: {:?} new temps, starting from index {}: {:?}",
 136                self.new_locals.len(), mir.local_decls.len(), self.new_locals);
 137         debug!("MirPatch: {} new blocks, starting from index {}",
 138                self.new_blocks.len(), mir.basic_blocks().len());
 139         mir.basic_blocks_mut().extend(self.new_blocks);
 140         mir.local_decls.extend(self.new_locals);
 141         for (src, patch) in self.patch_map.into_iter_enumerated() {
 142             if let Some(patch) = patch {
 143                 debug!("MirPatch: patching block {:?}", src);
 144                 mir[src].terminator_mut().kind = patch;
 145             }
 146         }
 147
 148         let mut new_statements = self.new_statements;
 149         new_statements.sort_by_key(|s| s.0);
 150
 151         let mut delta = 0;
 152         let mut last_bb = START_BLOCK;
 153         for (mut loc, stmt) in new_statements {
 154             if loc.block != last_bb {
 155                 delta = 0;
 156                 last_bb = loc.block;
 157             }
 158             debug!("MirPatch: adding statement {:?} at loc {:?}+{}",
 159                    stmt, loc, delta);
 160             loc.statement_index += delta;
 161             let source_info = Self::source_info_for_index(
 162                 &mir[loc.block], loc
 163             );
 164             mir[loc.block].statements.insert(
 165                 loc.statement_index, Statement {
 166                     source_info,
 167                     kind: stmt
 168                 });
 169             delta += 1;
 170         }
 171     }
 172
 173     pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
 174         match data.statements.get(loc.statement_index) {
 175             Some(stmt) => stmt.source_info,
 176             None => data.terminator().source_info
 177         }
 178     }
 179
 180     pub fn source_info_for_location(&self, mir: &Mir<'_>, loc: Location) -> SourceInfo {
 181         let data = match loc.block.index().checked_sub(mir.basic_blocks().len()) {
 182             Some(new) => &self.new_blocks[new],
 183             None => &mir[loc.block]
 184         };
 185         Self::source_info_for_index(data, loc)
 186     }
 187 }