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