src/librustc_data_structures/box_region.rs

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