proc_macro/bridge/
client.rs

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

proc_macro/bridge/
client.rs
