src/librustc_data_structures/box_region.rs

   1 use std::cell::Cell;
   2 use std::marker::PhantomData;
   3 use std::ops::{Generator, GeneratorState};
   4 use std::pin::Pin;
   5
   6 #[derive(Copy, Clone)]
   7 pub struct AccessAction(*mut dyn FnMut());
   8
   9 impl AccessAction {
  10     pub fn get(self) -> *mut dyn FnMut() {
  11         self.0
  12     }
  13 }
  14
  15 #[derive(Copy, Clone)]
  16 pub enum Action {
  17     Access(AccessAction),
  18     Complete,
  19 }
  20
  21 thread_local!(pub static BOX_REGION_ARG: Cell<Action> = Cell::new(Action::Complete));
  22
  23 pub struct PinnedGenerator<I, A, R> {
  24     generator: Pin<Box<dyn Generator<Yield = YieldType<I, A>, Return = R>>>,
  25 }
  26
  27 impl<I, A, R> PinnedGenerator<I, A, R> {
  28     #[cfg(bootstrap)]
  29     pub fn new<T: Generator<Yield = YieldType<I, A>, Return = R> + 'static>(
  30         generator: T,
  31     ) -> (I, Self) {
  32         let mut result = PinnedGenerator { generator: Box::pin(generator) };
  33
  34         // Run it to the first yield to set it up
  35         let init = match Pin::new(&mut result.generator).resume() {
  36             GeneratorState::Yielded(YieldType::Initial(y)) => y,
  37             _ => panic!(),
  38         };
  39
  40         (init, result)
  41     }
  42
  43     #[cfg(not(bootstrap))]
  44     pub fn new<T: Generator<Yield = YieldType<I, A>, Return = R> + 'static>(
  45         generator: T,
  46     ) -> (I, Self) {
  47         let mut result = PinnedGenerator { generator: Box::pin(generator) };
  48
  49         // Run it to the first yield to set it up
  50         let init = match Pin::new(&mut result.generator).resume(()) {
  51             GeneratorState::Yielded(YieldType::Initial(y)) => y,
  52             _ => panic!(),
  53         };
  54
  55         (init, result)
  56     }
  57
  58     #[cfg(bootstrap)]
  59     pub unsafe fn access(&mut self, closure: *mut dyn FnMut()) {
  60         BOX_REGION_ARG.with(|i| {
  61             i.set(Action::Access(AccessAction(closure)));
  62         });
  63
  64         // Call the generator, which in turn will call the closure in BOX_REGION_ARG
  65         if let GeneratorState::Complete(_) = Pin::new(&mut self.generator).resume() {
  66             panic!()
  67         }
  68     }
  69
  70     #[cfg(not(bootstrap))]
  71     pub unsafe fn access(&mut self, closure: *mut dyn FnMut()) {
  72         BOX_REGION_ARG.with(|i| {
  73             i.set(Action::Access(AccessAction(closure)));
  74         });
  75
  76         // Call the generator, which in turn will call the closure in BOX_REGION_ARG
  77         if let GeneratorState::Complete(_) = Pin::new(&mut self.generator).resume(()) {
  78             panic!()
  79         }
  80     }
  81
  82     #[cfg(bootstrap)]
  83     pub fn complete(&mut self) -> R {
  84         // Tell the generator we want it to complete, consuming it and yielding a result
  85         BOX_REGION_ARG.with(|i| i.set(Action::Complete));
  86
  87         let result = Pin::new(&mut self.generator).resume();
  88         if let GeneratorState::Complete(r) = result { r } else { panic!() }
  89     }
  90
  91     #[cfg(not(bootstrap))]
  92     pub fn complete(&mut self) -> R {
  93         // Tell the generator we want it to complete, consuming it and yielding a result
  94         BOX_REGION_ARG.with(|i| i.set(Action::Complete));
  95
  96         let result = Pin::new(&mut self.generator).resume(());
  97         if let GeneratorState::Complete(r) = result { r } else { panic!() }
  98     }
  99 }
 100
 101 #[derive(PartialEq)]
 102 pub struct Marker<T>(PhantomData<T>);
 103
 104 impl<T> Marker<T> {
 105     pub unsafe fn new() -> Self {
 106         Marker(PhantomData)
 107     }
 108 }
 109
 110 pub enum YieldType<I, A> {
 111     Initial(I),
 112     Accessor(Marker<A>),
 113 }
 114
 115 #[macro_export]
 116 #[allow_internal_unstable(fn_traits)]
 117 macro_rules! declare_box_region_type {
 118     (impl $v:vis
 119      $name: ident,
 120      $yield_type:ty,
 121      for($($lifetimes:tt)*),
 122      ($($args:ty),*) -> ($reti:ty, $retc:ty)
 123     ) => {
 124         $v struct $name($crate::box_region::PinnedGenerator<
 125             $reti,
 126             for<$($lifetimes)*> fn(($($args,)*)),
 127             $retc
 128         >);
 129
 130         impl $name {
 131             fn new<T: ::std::ops::Generator<Yield = $yield_type, Return = $retc> + 'static>(
 132                 generator: T
 133             ) -> ($reti, Self) {
 134                 let (initial, pinned) = $crate::box_region::PinnedGenerator::new(generator);
 135                 (initial, $name(pinned))
 136             }
 137
 138             $v fn access<F: for<$($lifetimes)*> FnOnce($($args,)*) -> R, R>(&mut self, f: F) -> R {
 139                 // Turn the FnOnce closure into *mut dyn FnMut()
 140                 // so we can pass it in to the generator using the BOX_REGION_ARG thread local
 141                 let mut r = None;
 142                 let mut f = Some(f);
 143                 let mut_f: &mut dyn for<$($lifetimes)*> FnMut(($($args,)*)) =
 144                     &mut |args| {
 145                         let f = f.take().unwrap();
 146                         r = Some(FnOnce::call_once(f, args));
 147                 };
 148                 let mut_f = mut_f as *mut dyn for<$($lifetimes)*> FnMut(($($args,)*));
 149
 150                 // Get the generator to call our closure
 151                 unsafe {
 152                     self.0.access(::std::mem::transmute(mut_f));
 153                 }
 154
 155                 // Unwrap the result
 156                 r.unwrap()
 157             }
 158
 159             $v fn complete(mut self) -> $retc {
 160                 self.0.complete()
 161             }
 162
 163             fn initial_yield(value: $reti) -> $yield_type {
 164                 $crate::box_region::YieldType::Initial(value)
 165             }
 166         }
 167     };
 168
 169     ($v:vis $name: ident, for($($lifetimes:tt)*), ($($args:ty),*) -> ($reti:ty, $retc:ty)) => {
 170         declare_box_region_type!(
 171             impl $v $name,
 172             $crate::box_region::YieldType<$reti, for<$($lifetimes)*> fn(($($args,)*))>,
 173             for($($lifetimes)*),
 174             ($($args),*) -> ($reti, $retc)
 175         );
 176     };
 177 }
 178
 179 #[macro_export]
 180 #[allow_internal_unstable(fn_traits)]
 181 macro_rules! box_region_allow_access {
 182     (for($($lifetimes:tt)*), ($($args:ty),*), ($($exprs:expr),*) ) => {
 183         loop {
 184             match $crate::box_region::BOX_REGION_ARG.with(|i| i.get()) {
 185                 $crate::box_region::Action::Access(accessor) => {
 186                     let accessor: &mut dyn for<$($lifetimes)*> FnMut($($args),*) = unsafe {
 187                         ::std::mem::transmute(accessor.get())
 188                     };
 189                     (*accessor)(($($exprs),*));
 190                     unsafe {
 191                         let marker = $crate::box_region::Marker::<
 192                             for<$($lifetimes)*> fn(($($args,)*))
 193                         >::new();
 194                         yield $crate::box_region::YieldType::Accessor(marker)
 195                     };
 196                 }
 197                 $crate::box_region::Action::Complete => break,
 198             }
 199         }
 200     }
 201 }