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