proc_macro\bridge/
client.rs

1//! Client-side types.
2
3use std::cell::RefCell;
4
5use super::*;
6
7pub(crate) struct TokenStream {
8    handle: handle::Handle,
9}
10
11impl !Send for TokenStream {}
12impl !Sync for TokenStream {}
13
14// Forward `Drop::drop` to the inherent `drop` method.
15impl Drop for TokenStream {
16    fn drop(&mut self) {
17        Methods::ts_drop(TokenStream { handle: self.handle });
18    }
19}
20
21impl<S> Encode<S> for TokenStream {
22    #[inline]
23    fn encode(self, w: &mut Buffer, s: &mut S) {
24        mem::ManuallyDrop::new(self).handle.encode(w, s);
25    }
26}
27
28impl<S> Encode<S> for &TokenStream {
29    #[inline]
30    fn encode(self, w: &mut Buffer, s: &mut S) {
31        self.handle.encode(w, s);
32    }
33}
34
35impl<S> Decode<'_, '_, S> for TokenStream {
36    #[inline]
37    fn decode(r: &mut &[u8], s: &mut S) -> Self {
38        TokenStream { handle: handle::Handle::decode(r, s) }
39    }
40}
41
42impl Encode<()> for crate::TokenStream {
43    #[inline]
44    fn encode(self, w: &mut Buffer, s: &mut ()) {
45        self.0.encode(w, s)
46    }
47}
48
49impl Decode<'_, '_, ()> for crate::TokenStream {
50    #[inline]
51    fn decode(r: &mut &[u8], s: &mut ()) -> Self {
52        crate::TokenStream(Some(Decode::decode(r, s)))
53    }
54}
55
56#[derive(Copy, Clone, PartialEq, Eq, Hash)]
57pub(crate) struct Span {
58    handle: handle::Handle,
59}
60
61impl !Send for Span {}
62impl !Sync for Span {}
63
64impl<S> Encode<S> for Span {
65    #[inline]
66    fn encode(self, w: &mut Buffer, s: &mut S) {
67        self.handle.encode(w, s);
68    }
69}
70
71impl<S> Decode<'_, '_, S> for Span {
72    #[inline]
73    fn decode(r: &mut &[u8], s: &mut S) -> Self {
74        Span { handle: handle::Handle::decode(r, s) }
75    }
76}
77
78impl Clone for TokenStream {
79    fn clone(&self) -> Self {
80        Methods::ts_clone(self)
81    }
82}
83
84impl Span {
85    pub(crate) fn def_site() -> Span {
86        Bridge::with(|bridge| bridge.globals.def_site)
87    }
88
89    pub(crate) fn call_site() -> Span {
90        Bridge::with(|bridge| bridge.globals.call_site)
91    }
92
93    pub(crate) fn mixed_site() -> Span {
94        Bridge::with(|bridge| bridge.globals.mixed_site)
95    }
96}
97
98impl fmt::Debug for Span {
99    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100        f.write_str(&Methods::span_debug(*self))
101    }
102}
103
104pub(crate) use super::Methods;
105pub(crate) use super::symbol::Symbol;
106
107macro_rules! define_client_side {
108    (
109        $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
110    ) => {
111        impl Methods {
112            $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)? {
113                Bridge::with(|bridge| {
114                    let mut buf = bridge.cached_buffer.take();
115
116                    buf.clear();
117                    ApiTags::$method.encode(&mut buf, &mut ());
118                    $($arg.encode(&mut buf, &mut ());)*
119
120                    buf = bridge.dispatch.call(buf);
121
122                    let r = Result::<_, PanicMessage>::decode(&mut &buf[..], &mut ());
123
124                    bridge.cached_buffer = buf;
125
126                    r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
127                })
128            })*
129        }
130    }
131}
132with_api!(define_client_side, TokenStream, Span, Symbol);
133
134struct Bridge<'a> {
135    /// Reusable buffer (only `clear`-ed, never shrunk), primarily
136    /// used for making requests.
137    cached_buffer: Buffer,
138
139    /// Server-side function that the client uses to make requests.
140    dispatch: closure::Closure<'a>,
141
142    /// Provided globals for this macro expansion.
143    globals: ExpnGlobals<Span>,
144}
145
146impl<'a> !Send for Bridge<'a> {}
147impl<'a> !Sync for Bridge<'a> {}
148
149#[allow(unsafe_code)]
150mod state {
151    use std::cell::{Cell, RefCell};
152    use std::ptr;
153
154    use super::Bridge;
155
156    thread_local! {
157        static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
158    }
159
160    pub(super) fn set<'bridge, R>(state: &RefCell<Bridge<'bridge>>, f: impl FnOnce() -> R) -> R {
161        struct RestoreOnDrop(*const ());
162        impl Drop for RestoreOnDrop {
163            fn drop(&mut self) {
164                BRIDGE_STATE.set(self.0);
165            }
166        }
167
168        let inner = ptr::from_ref(state).cast();
169        let outer = BRIDGE_STATE.replace(inner);
170        let _restore = RestoreOnDrop(outer);
171
172        f()
173    }
174
175    pub(super) fn with<R>(
176        f: impl for<'bridge> FnOnce(Option<&RefCell<Bridge<'bridge>>>) -> R,
177    ) -> R {
178        let state = BRIDGE_STATE.get();
179        // SAFETY: the only place where the pointer is set is in `set`. It puts
180        // back the previous value after the inner call has returned, so we know
181        // that as long as the pointer is not null, it came from a reference to
182        // a `RefCell<Bridge>` that outlasts the call to this function. Since `f`
183        // works the same for any lifetime of the bridge, including the actual
184        // one, we can lie here and say that the lifetime is `'static` without
185        // anyone noticing.
186        let bridge = unsafe { state.cast::<RefCell<Bridge<'static>>>().as_ref() };
187        f(bridge)
188    }
189}
190
191impl Bridge<'_> {
192    fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
193        state::with(|state| {
194            let bridge = state.expect("procedural macro API is used outside of a procedural macro");
195            let mut bridge = bridge
196                .try_borrow_mut()
197                .expect("procedural macro API is used while it's already in use");
198            f(&mut bridge)
199        })
200    }
201}
202
203pub(crate) fn is_available() -> bool {
204    state::with(|s| s.is_some())
205}
206
207/// A client-side RPC entry-point, which may be using a different `proc_macro`
208/// from the one used by the server, but can be invoked compatibly.
209///
210/// Note that the input and output type parameters are erased. They do not
211/// participate in the ABI, so while using the wrong runN method will likely
212/// result in a panic, it will not result in UB.
213#[repr(C)]
214#[derive(Copy, Clone)]
215pub struct Client {
216    pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
217}
218
219fn maybe_install_panic_hook(force_show_panics: bool) {
220    // Hide the default panic output within `proc_macro` expansions.
221    // NB. the server can't do this because it may use a different std.
222    static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
223    HIDE_PANICS_DURING_EXPANSION.call_once(|| {
224        let prev = panic::take_hook();
225        panic::set_hook(Box::new(move |info| {
226            // We normally report panics by catching unwinds and passing the payload from the
227            // unwind back to the compiler, but if the panic doesn't unwind we'll abort before
228            // the compiler has a chance to print an error. So we special-case PanicInfo where
229            // can_unwind is false.
230            if force_show_panics || !is_available() || !info.can_unwind() {
231                prev(info)
232            }
233        }));
234    });
235}
236
237/// Client-side helper for handling client panics, entering the bridge,
238/// deserializing input and serializing output.
239fn run_client<A: for<'a, 's> Decode<'a, 's, ()>>(
240    config: BridgeConfig<'_>,
241    f: impl FnOnce(A) -> crate::TokenStream,
242) -> Buffer {
243    let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
244
245    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
246        maybe_install_panic_hook(force_show_panics);
247
248        // Make sure the symbol store is empty before decoding inputs.
249        Symbol::invalidate_all();
250
251        let reader = &mut &buf[..];
252        let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
253
254        // Put the buffer we used for input back in the `Bridge` for requests.
255        let state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals });
256
257        let output = state::set(&state, || f(input));
258
259        // Take the `cached_buffer` back out, for the output value.
260        buf = RefCell::into_inner(state).cached_buffer;
261
262        output
263    }));
264
265    // Serialize response of type `Result<R, PanicMessage>`.
266    buf.clear();
267    res.map_err(PanicMessage::from).encode(&mut buf, &mut ());
268
269    // Now that a response has been serialized, invalidate all symbols
270    // registered with the interner.
271    Symbol::invalidate_all();
272    buf
273}
274
275impl Client {
276    pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
277        Client {
278            run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
279                run_client(bridge, |input| f(input))
280            }),
281        }
282    }
283
284    pub const fn expand2(
285        f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
286    ) -> Self {
287        Client {
288            run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
289                run_client(bridge, |(input, input2)| f(input, input2))
290            }),
291        }
292    }
293}
proc_macro\bridge/client.rs

proc_macro\bridge/
client.rs
