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