]>
Commit | Line | Data |
---|---|---|
532ac7d7 XL |
1 | use std::cell::Cell; |
2 | use std::marker::PhantomData; | |
532ac7d7 | 3 | use std::ops::{Generator, GeneratorState}; |
dfeec247 | 4 | use std::pin::Pin; |
532ac7d7 XL |
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> { | |
dfeec247 | 24 | generator: Pin<Box<dyn Generator<Yield = YieldType<I, A>, Return = R>>>, |
532ac7d7 XL |
25 | } |
26 | ||
27 | impl<I, A, R> PinnedGenerator<I, A, R> { | |
74b04a01 XL |
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 | ||
74b04a01 XL |
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 | ||
74b04a01 XL |
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 | } | |
532ac7d7 XL |
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 | } |