| #![feature(generators, generator_trait, never_type)] |
| |
| use std::ops::{GeneratorState::{self, *}, Generator}; |
| use std::pin::Pin; |
| use std::sync::atomic::{AtomicUsize, Ordering}; |
| use std::fmt::Debug; |
| use std::mem::ManuallyDrop; |
| use std::ptr; |
| |
| fn basic() { |
| fn finish<T>(mut amt: usize, mut t: T) -> T::Return |
| where T: Generator<Yield = usize> |
| { |
| // We are not moving the `t` around until it gets dropped, so this is okay. |
| let mut t = unsafe { Pin::new_unchecked(&mut t) }; |
| loop { |
| let state = t.as_mut().resume(()); |
| // Test if the generator is valid (according to type invariants). |
| let _ = unsafe { ManuallyDrop::new(ptr::read(t.as_mut().get_unchecked_mut())) }; |
| match state { |
| GeneratorState::Yielded(y) => { |
| amt -= y; |
| } |
| GeneratorState::Complete(ret) => { |
| assert_eq!(amt, 0); |
| return ret |
| } |
| } |
| } |
| } |
| |
| enum Never {} |
| fn never() -> Never { |
| panic!() |
| } |
| |
| finish(1, || yield 1); |
| |
| finish(3, || { |
| let mut x = 0; |
| yield 1; |
| x += 1; |
| yield 1; |
| x += 1; |
| yield 1; |
| assert_eq!(x, 2); |
| }); |
| |
| finish(7*8/2, || { |
| for i in 0..8 { |
| yield i; |
| } |
| }); |
| |
| finish(1, || { |
| if true { |
| yield 1; |
| } else { |
| } |
| }); |
| |
| finish(1, || { |
| if false { |
| } else { |
| yield 1; |
| } |
| }); |
| |
| finish(2, || { |
| if { yield 1; false } { |
| yield 1; |
| panic!() |
| } |
| yield 1; |
| }); |
| |
| // also test a self-referential generator |
| assert_eq!( |
| finish(5, || { |
| let mut x = Box::new(5); |
| let y = &mut *x; |
| *y = 5; |
| yield *y; |
| *y = 10; |
| *x |
| }), |
| 10 |
| ); |
| |
| let b = true; |
| finish(1, || { |
| yield 1; |
| if b { return; } |
| #[allow(unused)] |
| let x = never(); |
| yield 2; |
| drop(x); |
| }); |
| |
| finish(3, || { |
| yield 1; |
| #[allow(unreachable_code)] |
| let _x: (String, !) = (String::new(), { yield 2; return }); |
| }); |
| } |
| |
| fn smoke_resume_arg() { |
| fn drain<G: Generator<R, Yield = Y> + Unpin, R, Y>( |
| gen: &mut G, |
| inout: Vec<(R, GeneratorState<Y, G::Return>)>, |
| ) where |
| Y: Debug + PartialEq, |
| G::Return: Debug + PartialEq, |
| { |
| let mut gen = Pin::new(gen); |
| |
| for (input, out) in inout { |
| assert_eq!(gen.as_mut().resume(input), out); |
| // Test if the generator is valid (according to type invariants). |
| let _ = unsafe { ManuallyDrop::new(ptr::read(gen.as_mut().get_unchecked_mut())) }; |
| } |
| } |
| |
| static DROPS: AtomicUsize = AtomicUsize::new(0); |
| |
| #[derive(Debug, PartialEq)] |
| struct DropMe; |
| |
| impl Drop for DropMe { |
| fn drop(&mut self) { |
| DROPS.fetch_add(1, Ordering::SeqCst); |
| } |
| } |
| |
| fn expect_drops<T>(expected_drops: usize, f: impl FnOnce() -> T) -> T { |
| DROPS.store(0, Ordering::SeqCst); |
| |
| let res = f(); |
| |
| let actual_drops = DROPS.load(Ordering::SeqCst); |
| assert_eq!(actual_drops, expected_drops); |
| res |
| } |
| |
| drain( |
| &mut |mut b| { |
| while b != 0 { |
| b = yield (b + 1); |
| } |
| -1 |
| }, |
| vec![(1, Yielded(2)), (-45, Yielded(-44)), (500, Yielded(501)), (0, Complete(-1))], |
| ); |
| |
| expect_drops(2, || drain(&mut |a| yield a, vec![(DropMe, Yielded(DropMe))])); |
| |
| expect_drops(6, || { |
| drain( |
| &mut |a| yield yield a, |
| vec![(DropMe, Yielded(DropMe)), (DropMe, Yielded(DropMe)), (DropMe, Complete(DropMe))], |
| ) |
| }); |
| |
| #[allow(unreachable_code)] |
| expect_drops(2, || drain(&mut |a| yield return a, vec![(DropMe, Complete(DropMe))])); |
| |
| expect_drops(2, || { |
| drain( |
| &mut |a: DropMe| { |
| if false { yield () } else { a } |
| }, |
| vec![(DropMe, Complete(DropMe))], |
| ) |
| }); |
| |
| expect_drops(4, || { |
| drain( |
| #[allow(unused_assignments, unused_variables)] |
| &mut |mut a: DropMe| { |
| a = yield; |
| a = yield; |
| a = yield; |
| }, |
| vec![ |
| (DropMe, Yielded(())), |
| (DropMe, Yielded(())), |
| (DropMe, Yielded(())), |
| (DropMe, Complete(())), |
| ], |
| ) |
| }); |
| } |
| |
| fn main() { |
| basic(); |
| smoke_resume_arg(); |
| } |