Make generator transform move resume arg around

The resume arg is passed as argument `_2` and needs to be moved to the
`Yield`s target `Place`
This commit is contained in:
Jonas Schievink 2020-01-25 02:33:52 +01:00
parent 3c069a066e
commit 3c22e51e7f

View File

@ -192,9 +192,10 @@ const RETURNED: usize = GeneratorSubsts::RETURNED;
/// Generator has been poisoned /// Generator has been poisoned
const POISONED: usize = GeneratorSubsts::POISONED; const POISONED: usize = GeneratorSubsts::POISONED;
struct SuspensionPoint { struct SuspensionPoint<'tcx> {
state: usize, state: usize,
resume: BasicBlock, resume: BasicBlock,
resume_arg: Place<'tcx>,
drop: Option<BasicBlock>, drop: Option<BasicBlock>,
storage_liveness: liveness::LiveVarSet, storage_liveness: liveness::LiveVarSet,
} }
@ -216,7 +217,7 @@ struct TransformVisitor<'tcx> {
storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>, storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
// A list of suspension points, generated during the transform // A list of suspension points, generated during the transform
suspension_points: Vec<SuspensionPoint>, suspension_points: Vec<SuspensionPoint<'tcx>>,
// The original RETURN_PLACE local // The original RETURN_PLACE local
new_ret_local: Local, new_ret_local: Local,
@ -303,8 +304,8 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
Operand::Move(Place::from(self.new_ret_local)), Operand::Move(Place::from(self.new_ret_local)),
None, None,
)), )),
TerminatorKind::Yield { ref value, resume, drop } => { TerminatorKind::Yield { ref value, resume, resume_arg, drop } => {
Some((VariantIdx::new(0), Some(resume), value.clone(), drop)) Some((VariantIdx::new(0), Some((resume, resume_arg)), value.clone(), drop))
} }
_ => None, _ => None,
}; };
@ -319,13 +320,14 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
self.make_state(state_idx, v), self.make_state(state_idx, v),
)), )),
}); });
let state = if let Some(resume) = resume { let state = if let Some((resume, resume_arg)) = resume {
// Yield // Yield
let state = 3 + self.suspension_points.len(); let state = 3 + self.suspension_points.len();
self.suspension_points.push(SuspensionPoint { self.suspension_points.push(SuspensionPoint {
state, state,
resume, resume,
resume_arg,
drop, drop,
storage_liveness: self.storage_liveness.get(&block).unwrap().clone(), storage_liveness: self.storage_liveness.get(&block).unwrap().clone(),
}); });
@ -1063,7 +1065,7 @@ fn create_cases<'tcx, F>(
target: F, target: F,
) -> Vec<(usize, BasicBlock)> ) -> Vec<(usize, BasicBlock)>
where where
F: Fn(&SuspensionPoint) -> Option<BasicBlock>, F: Fn(&SuspensionPoint<'tcx>) -> Option<BasicBlock>,
{ {
let source_info = source_info(body); let source_info = source_info(body);
@ -1085,6 +1087,16 @@ where
} }
} }
// Move the resume argument to the destination place of the `Yield` terminator
let resume_arg = Local::new(2); // 0 = return, 1 = self
statements.push(Statement {
source_info,
kind: StatementKind::Assign(box (
point.resume_arg,
Rvalue::Use(Operand::Move(resume_arg.into())),
)),
});
// Then jump to the real target // Then jump to the real target
body.basic_blocks_mut().push(BasicBlockData { body.basic_blocks_mut().push(BasicBlockData {
statements, statements,
@ -1163,7 +1175,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
}; };
transform.visit_body(body); transform.visit_body(body);
// Update our MIR struct to reflect the changed we've made // Update our MIR struct to reflect the changes we've made
body.yield_ty = None; body.yield_ty = None;
body.arg_count = 2; // self, resume arg body.arg_count = 2; // self, resume arg
body.spread_arg = None; body.spread_arg = None;