src/librustc_mir/util/patch.rs

   1 // Copyright 2016 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use rustc::ty::Ty;
  12 use rustc::mir::*;
  13 use rustc_data_structures::indexed_vec::{IndexVec, Idx};
  14 use syntax_pos::Span;
  15
  16 /// This struct represents a patch to MIR, which can add
  17 /// new statements and basic blocks and patch over block
  18 /// terminators.
  19 pub struct MirPatch<'tcx> {
  20     patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
  21     new_blocks: Vec<BasicBlockData<'tcx>>,
  22     new_statements: Vec<(Location, StatementKind<'tcx>)>,
  23     new_locals: Vec<LocalDecl<'tcx>>,
  24     resume_block: BasicBlock,
  25     next_local: usize,
  26     make_nop: Vec<Location>,
  27 }
  28
  29 impl<'tcx> MirPatch<'tcx> {
  30     pub fn new(mir: &Mir<'tcx>) -> Self {
  31         let mut result = MirPatch {
  32             patch_map: IndexVec::from_elem(None, mir.basic_blocks()),
  33             new_blocks: vec![],
  34             new_statements: vec![],
  35             new_locals: vec![],
  36             next_local: mir.local_decls.len(),
  37             resume_block: START_BLOCK,
  38             make_nop: vec![]
  39         };
  40
  41         // make sure the MIR we create has a resume block. It is
  42         // completely legal to convert jumps to the resume block
  43         // to jumps to None, but we occasionally have to add
  44         // instructions just before that.
  45
  46         let mut resume_block = None;
  47         let mut resume_stmt_block = None;
  48         for (bb, block) in mir.basic_blocks().iter_enumerated() {
  49             if let TerminatorKind::Resume = block.terminator().kind {
  50                 if block.statements.len() > 0 {
  51                     assert!(resume_stmt_block.is_none());
  52                     resume_stmt_block = Some(bb);
  53                 } else {
  54                     resume_block = Some(bb);
  55                 }
  56                 break
  57             }
  58         }
  59         let resume_block = resume_block.unwrap_or_else(|| {
  60             result.new_block(BasicBlockData {
  61                 statements: vec![],
  62                 terminator: Some(Terminator {
  63                     source_info: SourceInfo {
  64                         span: mir.span,
  65                         scope: OUTERMOST_SOURCE_SCOPE
  66                     },
  67                     kind: TerminatorKind::Resume
  68                 }),
  69                 is_cleanup: true
  70             })});
  71         result.resume_block = resume_block;
  72         if let Some(resume_stmt_block) = resume_stmt_block {
  73             result.patch_terminator(resume_stmt_block, TerminatorKind::Goto {
  74                 target: resume_block
  75             });
  76         }
  77         result
  78     }
  79
  80     pub fn resume_block(&self) -> BasicBlock {
  81         self.resume_block
  82     }
  83
  84     pub fn is_patched(&self, bb: BasicBlock) -> bool {
  85         self.patch_map[bb].is_some()
  86     }
  87
  88     pub fn terminator_loc(&self, mir: &Mir<'tcx>, bb: BasicBlock) -> Location {
  89         let offset = match bb.index().checked_sub(mir.basic_blocks().len()) {
  90             Some(index) => self.new_blocks[index].statements.len(),
  91             None => mir[bb].statements.len()
  92         };
  93         Location {
  94             block: bb,
  95             statement_index: offset
  96         }
  97     }
  98
  99     pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
 100         let index = self.next_local;
 101         self.next_local += 1;
 102         self.new_locals.push(LocalDecl::new_temp(ty, span));
 103         Local::new(index as usize)
 104     }
 105
 106     pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
 107         let index = self.next_local;
 108         self.next_local += 1;
 109         self.new_locals.push(LocalDecl::new_internal(ty, span));
 110         Local::new(index as usize)
 111     }
 112
 113     pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
 114         let block = BasicBlock::new(self.patch_map.len());
 115         debug!("MirPatch: new_block: {:?}: {:?}", block, data);
 116         self.new_blocks.push(data);
 117         self.patch_map.push(None);
 118         block
 119     }
 120
 121     pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
 122         assert!(self.patch_map[block].is_none());
 123         debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
 124         self.patch_map[block] = Some(new);
 125     }
 126
 127     pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
 128         debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
 129         self.new_statements.push((loc, stmt));
 130     }
 131
 132     pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
 133         self.add_statement(loc, StatementKind::Assign(place, rv));
 134     }
 135
 136     pub fn make_nop(&mut self, loc: Location) {
 137         self.make_nop.push(loc);
 138     }
 139
 140     pub fn apply(self, mir: &mut Mir<'tcx>) {
 141         debug!("MirPatch: make nops at: {:?}", self.make_nop);
 142         for loc in self.make_nop {
 143             mir.make_statement_nop(loc);
 144         }
 145         debug!("MirPatch: {:?} new temps, starting from index {}: {:?}",
 146                self.new_locals.len(), mir.local_decls.len(), self.new_locals);
 147         debug!("MirPatch: {} new blocks, starting from index {}",
 148                self.new_blocks.len(), mir.basic_blocks().len());
 149         mir.basic_blocks_mut().extend(self.new_blocks);
 150         mir.local_decls.extend(self.new_locals);
 151         for (src, patch) in self.patch_map.into_iter_enumerated() {
 152             if let Some(patch) = patch {
 153                 debug!("MirPatch: patching block {:?}", src);
 154                 mir[src].terminator_mut().kind = patch;
 155             }
 156         }
 157
 158         let mut new_statements = self.new_statements;
 159         new_statements.sort_by(|u,v| u.0.cmp(&v.0));
 160
 161         let mut delta = 0;
 162         let mut last_bb = START_BLOCK;
 163         for (mut loc, stmt) in new_statements {
 164             if loc.block != last_bb {
 165                 delta = 0;
 166                 last_bb = loc.block;
 167             }
 168             debug!("MirPatch: adding statement {:?} at loc {:?}+{}",
 169                    stmt, loc, delta);
 170             loc.statement_index += delta;
 171             let source_info = Self::source_info_for_index(
 172                 &mir[loc.block], loc
 173             );
 174             mir[loc.block].statements.insert(
 175                 loc.statement_index, Statement {
 176                     source_info,
 177                     kind: stmt
 178                 });
 179             delta += 1;
 180         }
 181     }
 182
 183     pub fn source_info_for_index(data: &BasicBlockData, loc: Location) -> SourceInfo {
 184         match data.statements.get(loc.statement_index) {
 185             Some(stmt) => stmt.source_info,
 186             None => data.terminator().source_info
 187         }
 188     }
 189
 190     pub fn source_info_for_location(&self, mir: &Mir, loc: Location) -> SourceInfo {
 191         let data = match loc.block.index().checked_sub(mir.basic_blocks().len()) {
 192             Some(new) => &self.new_blocks[new],
 193             None => &mir[loc.block]
 194         };
 195         Self::source_info_for_index(data, loc)
 196     }
 197 }