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