src/tools/miri/tests/pass/generator.rs

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