compiler/rustc_middle/src/mir/patch.rs

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