1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use std::cell::Cell;
use std::marker::PhantomData;
use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

#[derive(Copy, Clone)]
pub struct AccessAction(*mut dyn FnMut());

impl AccessAction {
    pub fn get(self) -> *mut dyn FnMut() {
        self.0
    }
}

#[derive(Copy, Clone)]
pub enum Action {
    Access(AccessAction),
    Complete,
}

thread_local!(pub static BOX_REGION_ARG: Cell<Action> = Cell::new(Action::Complete));

pub struct PinnedGenerator<I, A, R> {
    generator: Pin<Box<dyn Generator<Yield = YieldType<I, A>, Return = R>>>,
}

impl<I, A, R> PinnedGenerator<I, A, R> {
    pub fn new<T: Generator<Yield = YieldType<I, A>, Return = R> + 'static>(
        generator: T,
    ) -> (I, Self) {
        let mut result = PinnedGenerator { generator: Box::pin(generator) };

        // Run it to the first yield to set it up
        let init = match Pin::new(&mut result.generator).resume() {
            GeneratorState::Yielded(YieldType::Initial(y)) => y,
            _ => panic!(),
        };

        (init, result)
    }

    pub unsafe fn access(&mut self, closure: *mut dyn FnMut()) {
        BOX_REGION_ARG.with(|i| {
            i.set(Action::Access(AccessAction(closure)));
        });

        // Call the generator, which in turn will call the closure in BOX_REGION_ARG
        if let GeneratorState::Complete(_) = Pin::new(&mut self.generator).resume() {
            panic!()
        }
    }

    pub fn complete(&mut self) -> R {
        // Tell the generator we want it to complete, consuming it and yielding a result
        BOX_REGION_ARG.with(|i| i.set(Action::Complete));

        let result = Pin::new(&mut self.generator).resume();
        if let GeneratorState::Complete(r) = result { r } else { panic!() }
    }
}

#[derive(PartialEq)]
pub struct Marker<T>(PhantomData<T>);

impl<T> Marker<T> {
    pub unsafe fn new() -> Self {
        Marker(PhantomData)
    }
}

pub enum YieldType<I, A> {
    Initial(I),
    Accessor(Marker<A>),
}

#[macro_export]
#[allow_internal_unstable(fn_traits)]
macro_rules! declare_box_region_type {
    (impl $v:vis
     $name: ident,
     $yield_type:ty,
     for($($lifetimes:tt)*),
     ($($args:ty),*) -> ($reti:ty, $retc:ty)
    ) => {
        $v struct $name($crate::box_region::PinnedGenerator<
            $reti,
            for<$($lifetimes)*> fn(($($args,)*)),
            $retc
        >);

        impl $name {
            fn new<T: ::std::ops::Generator<Yield = $yield_type, Return = $retc> + 'static>(
                generator: T
            ) -> ($reti, Self) {
                let (initial, pinned) = $crate::box_region::PinnedGenerator::new(generator);
                (initial, $name(pinned))
            }

            $v fn access<F: for<$($lifetimes)*> FnOnce($($args,)*) -> R, R>(&mut self, f: F) -> R {
                // Turn the FnOnce closure into *mut dyn FnMut()
                // so we can pass it in to the generator using the BOX_REGION_ARG thread local
                let mut r = None;
                let mut f = Some(f);
                let mut_f: &mut dyn for<$($lifetimes)*> FnMut(($($args,)*)) =
                    &mut |args| {
                        let f = f.take().unwrap();
                        r = Some(FnOnce::call_once(f, args));
                };
                let mut_f = mut_f as *mut dyn for<$($lifetimes)*> FnMut(($($args,)*));

                // Get the generator to call our closure
                unsafe {
                    self.0.access(::std::mem::transmute(mut_f));
                }

                // Unwrap the result
                r.unwrap()
            }

            $v fn complete(mut self) -> $retc {
                self.0.complete()
            }

            fn initial_yield(value: $reti) -> $yield_type {
                $crate::box_region::YieldType::Initial(value)
            }
        }
    };

    ($v:vis $name: ident, for($($lifetimes:tt)*), ($($args:ty),*) -> ($reti:ty, $retc:ty)) => {
        declare_box_region_type!(
            impl $v $name,
            $crate::box_region::YieldType<$reti, for<$($lifetimes)*> fn(($($args,)*))>,
            for($($lifetimes)*),
            ($($args),*) -> ($reti, $retc)
        );
    };
}

#[macro_export]
#[allow_internal_unstable(fn_traits)]
macro_rules! box_region_allow_access {
    (for($($lifetimes:tt)*), ($($args:ty),*), ($($exprs:expr),*) ) => {
        loop {
            match $crate::box_region::BOX_REGION_ARG.with(|i| i.get()) {
                $crate::box_region::Action::Access(accessor) => {
                    let accessor: &mut dyn for<$($lifetimes)*> FnMut($($args),*) = unsafe {
                        ::std::mem::transmute(accessor.get())
                    };
                    (*accessor)(($($exprs),*));
                    unsafe {
                        let marker = $crate::box_region::Marker::<
                            for<$($lifetimes)*> fn(($($args,)*))
                        >::new();
                        yield $crate::box_region::YieldType::Accessor(marker)
                    };
                }
                $crate::box_region::Action::Complete => break,
            }
        }
    }
}