tests/run-pass/generator.rs

   1 #![feature(generators, generator_trait, never_type)]
   2
   3 use std::ops::{GeneratorState::{self, *}, Generator};
   4 use std::pin::Pin;
   5 use std::sync::atomic::{AtomicUsize, Ordering};
   6 use std::fmt::Debug;
   7 use std::mem::ManuallyDrop;
   8 use std::ptr;
   9
  10 fn basic() {
  11     fn finish<T>(mut amt: usize, mut t: T) -> T::Return
  12         where T: Generator<Yield = usize>
  13     {
  14         // We are not moving the `t` around until it gets dropped, so this is okay.
  15         let mut t = unsafe { Pin::new_unchecked(&mut t) };
  16         loop {
  17             let state = t.as_mut().resume(());
  18             // Test if the generator is valid (according to type invariants).
  19             let _ = unsafe { ManuallyDrop::new(ptr::read(t.as_mut().get_unchecked_mut())) };
  20             match state {
  21                 GeneratorState::Yielded(y) => {
  22                     amt -= y;
  23                 }
  24                 GeneratorState::Complete(ret) => {
  25                     assert_eq!(amt, 0);
  26                     return ret
  27                 }
  28             }
  29         }
  30     }
  31
  32     enum Never {}
  33     fn never() -> Never {
  34         panic!()
  35     }
  36
  37     finish(1, || yield 1);
  38
  39     finish(3, || {
  40         let mut x = 0;
  41         yield 1;
  42         x += 1;
  43         yield 1;
  44         x += 1;
  45         yield 1;
  46         assert_eq!(x, 2);
  47     });
  48
  49     finish(7*8/2, || {
  50         for i in 0..8 {
  51             yield i;
  52         }
  53     });
  54
  55     finish(1, || {
  56         if true {
  57             yield 1;
  58         } else {
  59         }
  60     });
  61
  62     finish(1, || {
  63         if false {
  64         } else {
  65             yield 1;
  66         }
  67     });
  68
  69     finish(2, || {
  70         if { yield 1; false } {
  71             yield 1;
  72             panic!()
  73         }
  74         yield 1;
  75     });
  76
  77     // also test a self-referential generator
  78     assert_eq!(
  79         finish(5, || {
  80             let mut x = Box::new(5);
  81             let y = &mut *x;
  82             *y = 5;
  83             yield *y;
  84             *y = 10;
  85             *x
  86         }),
  87         10
  88     );
  89
  90     let b = true;
  91     finish(1, || {
  92         yield 1;
  93         if b { return; }
  94         #[allow(unused)]
  95         let x = never();
  96         yield 2;
  97         drop(x);
  98     });
  99
 100     finish(3, || {
 101         yield 1;
 102         #[allow(unreachable_code)]
 103         let _x: (String, !) = (String::new(), { yield 2; return });
 104     });
 105 }
 106
 107 fn smoke_resume_arg() {
 108     fn drain<G: Generator<R, Yield = Y> + Unpin, R, Y>(
 109         gen: &mut G,
 110         inout: Vec<(R, GeneratorState<Y, G::Return>)>,
 111     ) where
 112         Y: Debug + PartialEq,
 113         G::Return: Debug + PartialEq,
 114     {
 115         let mut gen = Pin::new(gen);
 116
 117         for (input, out) in inout {
 118             assert_eq!(gen.as_mut().resume(input), out);
 119             // Test if the generator is valid (according to type invariants).
 120             let _ = unsafe { ManuallyDrop::new(ptr::read(gen.as_mut().get_unchecked_mut())) };
 121         }
 122     }
 123
 124     static DROPS: AtomicUsize = AtomicUsize::new(0);
 125
 126     #[derive(Debug, PartialEq)]
 127     struct DropMe;
 128
 129     impl Drop for DropMe {
 130         fn drop(&mut self) {
 131             DROPS.fetch_add(1, Ordering::SeqCst);
 132         }
 133     }
 134
 135     fn expect_drops<T>(expected_drops: usize, f: impl FnOnce() -> T) -> T {
 136         DROPS.store(0, Ordering::SeqCst);
 137
 138         let res = f();
 139
 140         let actual_drops = DROPS.load(Ordering::SeqCst);
 141         assert_eq!(actual_drops, expected_drops);
 142         res
 143     }
 144
 145     drain(
 146         &mut |mut b| {
 147             while b != 0 {
 148                 b = yield (b + 1);
 149             }
 150             -1
 151         },
 152         vec![(1, Yielded(2)), (-45, Yielded(-44)), (500, Yielded(501)), (0, Complete(-1))],
 153     );
 154
 155     expect_drops(2, || drain(&mut |a| yield a, vec![(DropMe, Yielded(DropMe))]));
 156
 157     expect_drops(6, || {
 158         drain(
 159             &mut |a| yield yield a,
 160             vec![(DropMe, Yielded(DropMe)), (DropMe, Yielded(DropMe)), (DropMe, Complete(DropMe))],
 161         )
 162     });
 163
 164     #[allow(unreachable_code)]
 165     expect_drops(2, || drain(&mut |a| yield return a, vec![(DropMe, Complete(DropMe))]));
 166
 167     expect_drops(2, || {
 168         drain(
 169             &mut |a: DropMe| {
 170                 if false { yield () } else { a }
 171             },
 172             vec![(DropMe, Complete(DropMe))],
 173         )
 174     });
 175
 176     expect_drops(4, || {
 177         drain(
 178             #[allow(unused_assignments, unused_variables)]
 179             &mut |mut a: DropMe| {
 180                 a = yield;
 181                 a = yield;
 182                 a = yield;
 183             },
 184             vec![
 185                 (DropMe, Yielded(())),
 186                 (DropMe, Yielded(())),
 187                 (DropMe, Yielded(())),
 188                 (DropMe, Complete(())),
 189             ],
 190         )
 191     });
 192 }
 193
 194 fn main() {
 195     basic();
 196     smoke_resume_arg();
 197 }