From 392e59500a96be718383e127d38bf74300f521c0 Mon Sep 17 00:00:00 2001 From: Jonas Schievink Date: Sun, 2 Feb 2020 22:46:06 +0100 Subject: [PATCH] Fix miscompilation --- src/librustc_mir/transform/generator.rs | 63 ++++++++++++++++--- src/test/mir-opt/generator-drop-cleanup.rs | 5 +- .../ui/generator/resume-live-across-yield.rs | 45 +++++++++++++ 3 files changed, 100 insertions(+), 13 deletions(-) create mode 100644 src/test/ui/generator/resume-live-across-yield.rs diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs index b5edcbe457b..414d91a18be 100644 --- a/src/librustc_mir/transform/generator.rs +++ b/src/librustc_mir/transform/generator.rs @@ -380,28 +380,35 @@ fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body PinArgVisitor { ref_gen_ty, tcx }.visit_body(body); } -fn replace_result_variable<'tcx>( - ret_ty: Ty<'tcx>, +/// Allocates a new local and replaces all references of `local` with it. Returns the new local. +/// +/// `local` will be changed to a new local decl with type `ty`. +/// +/// Note that the new local will be uninitialized. It is the caller's responsibility to assign some +/// valid value to it before its first use. +fn replace_local<'tcx>( + local: Local, + ty: Ty<'tcx>, body: &mut BodyAndCache<'tcx>, tcx: TyCtxt<'tcx>, ) -> Local { let source_info = source_info(body); - let new_ret = LocalDecl { + let new_decl = LocalDecl { mutability: Mutability::Mut, - ty: ret_ty, + ty, user_ty: UserTypeProjections::none(), source_info, internal: false, is_block_tail: None, local_info: LocalInfo::Other, }; - let new_ret_local = Local::new(body.local_decls.len()); - body.local_decls.push(new_ret); - body.local_decls.swap(RETURN_PLACE, new_ret_local); + let new_local = Local::new(body.local_decls.len()); + body.local_decls.push(new_decl); + body.local_decls.swap(local, new_local); - RenameLocalVisitor { from: RETURN_PLACE, to: new_ret_local, tcx }.visit_body(body); + RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body); - new_ret_local + new_local } struct StorageIgnored(liveness::LiveVarSet); @@ -794,6 +801,10 @@ fn compute_layout<'tcx>( (remap, layout, storage_liveness) } +/// Replaces the entry point of `body` with a block that switches on the generator discriminant and +/// dispatches to blocks according to `cases`. +/// +/// After this function, the former entry point of the function will be bb1. fn insert_switch<'tcx>( body: &mut BodyAndCache<'tcx>, cases: Vec<(usize, BasicBlock)>, @@ -1093,6 +1104,13 @@ fn create_cases<'tcx>( // Create StorageLive instructions for locals with live storage for i in 0..(body.local_decls.len()) { + if i == 2 { + // The resume argument is live on function entry. Don't insert a + // `StorageLive`, or the following `Assign` will read from uninitialized + // memory. + continue; + } + let l = Local::new(i); if point.storage_liveness.contains(l) && !transform.remap.contains_key(&l) { statements @@ -1110,6 +1128,10 @@ fn create_cases<'tcx>( Rvalue::Use(Operand::Move(resume_arg.into())), )), }); + statements.push(Statement { + source_info, + kind: StatementKind::StorageDead(resume_arg), + }); } // Then jump to the real target @@ -1166,7 +1188,28 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyAn // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local // RETURN_PLACE then is a fresh unused local with type ret_ty. - let new_ret_local = replace_result_variable(ret_ty, body, tcx); + let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx); + + // We also replace the resume argument and insert an `Assign`. + // This is needed because the resume argument might be live across a `yield`, and the + // transform assumes that any local live across a `yield` is assigned to before that. + let resume_local = Local::new(2); + let new_resume_local = + replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx); + + // When first entering the generator, move the resume argument into its new local. + let source_info = source_info(body); + let stmts = &mut body.basic_blocks_mut()[BasicBlock::new(0)].statements; + stmts.insert( + 0, + Statement { + source_info, + kind: StatementKind::Assign(box ( + new_resume_local.into(), + Rvalue::Use(Operand::Move(resume_local.into())), + )), + }, + ); // Extract locals which are live across suspension point into `layout` // `remap` gives a mapping from local indices onto generator struct indices diff --git a/src/test/mir-opt/generator-drop-cleanup.rs b/src/test/mir-opt/generator-drop-cleanup.rs index e5f3f9221c0..278dc49c926 100644 --- a/src/test/mir-opt/generator-drop-cleanup.rs +++ b/src/test/mir-opt/generator-drop-cleanup.rs @@ -13,8 +13,8 @@ fn main() { // START rustc.main-{{closure}}.generator_drop.0.mir // bb0: { -// _6 = discriminant((*_1)); -// switchInt(move _6) -> [0u32: bb4, 3u32: bb7, otherwise: bb8]; +// _7 = discriminant((*_1)); +// switchInt(move _7) -> [0u32: bb4, 3u32: bb7, otherwise: bb8]; // } // bb1: { // StorageDead(_4); @@ -37,7 +37,6 @@ fn main() { // goto -> bb3; // } // bb7: { -// StorageLive(_2); // StorageLive(_3); // StorageLive(_4); // goto -> bb1; diff --git a/src/test/ui/generator/resume-live-across-yield.rs b/src/test/ui/generator/resume-live-across-yield.rs new file mode 100644 index 00000000000..4c4cf117a55 --- /dev/null +++ b/src/test/ui/generator/resume-live-across-yield.rs @@ -0,0 +1,45 @@ +// run-pass + +#![feature(generators, generator_trait)] + +use std::ops::{Generator, GeneratorState}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; + +static DROP: AtomicUsize = AtomicUsize::new(0); + +#[derive(PartialEq, Eq, Debug)] +struct Dropper(String); + +impl Drop for Dropper { + fn drop(&mut self) { + DROP.fetch_add(1, Ordering::SeqCst); + } +} + +fn main() { + let mut g = |mut _d| { + _d = yield; + _d + }; + + let mut g = Pin::new(&mut g); + + assert_eq!( + g.as_mut().resume(Dropper(String::from("Hello world!"))), + GeneratorState::Yielded(()) + ); + assert_eq!(DROP.load(Ordering::Acquire), 0); + match g.as_mut().resume(Dropper(String::from("Number Two"))) { + GeneratorState::Complete(dropper) => { + assert_eq!(DROP.load(Ordering::Acquire), 1); + assert_eq!(dropper.0, "Number Two"); + drop(dropper); + assert_eq!(DROP.load(Ordering::Acquire), 2); + } + _ => unreachable!(), + } + + drop(g); + assert_eq!(DROP.load(Ordering::Acquire), 2); +} -- 2.44.0