library/proc_macro/src/bridge/client.rs

   1 //! Client-side types.
   2
   3 use super::*;
   4
   5 use std::marker::PhantomData;
   6
   7 macro_rules! define_handles {
   8     (
   9         'owned: $($oty:ident,)*
  10         'interned: $($ity:ident,)*
  11     ) => {
  12         #[repr(C)]
  13         #[allow(non_snake_case)]
  14         pub struct HandleCounters {
  15             $($oty: AtomicUsize,)*
  16             $($ity: AtomicUsize,)*
  17         }
  18
  19         impl HandleCounters {
  20             // FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
  21             // a wrapper `fn` pointer, once `const fn` can reference `static`s.
  22             extern "C" fn get() -> &'static Self {
  23                 static COUNTERS: HandleCounters = HandleCounters {
  24                     $($oty: AtomicUsize::new(1),)*
  25                     $($ity: AtomicUsize::new(1),)*
  26                 };
  27                 &COUNTERS
  28             }
  29         }
  30
  31         // FIXME(eddyb) generate the definition of `HandleStore` in `server.rs`.
  32         #[allow(non_snake_case)]
  33         pub(super) struct HandleStore<S: server::Types> {
  34             $($oty: handle::OwnedStore<S::$oty>,)*
  35             $($ity: handle::InternedStore<S::$ity>,)*
  36         }
  37
  38         impl<S: server::Types> HandleStore<S> {
  39             pub(super) fn new(handle_counters: &'static HandleCounters) -> Self {
  40                 HandleStore {
  41                     $($oty: handle::OwnedStore::new(&handle_counters.$oty),)*
  42                     $($ity: handle::InternedStore::new(&handle_counters.$ity),)*
  43                 }
  44             }
  45         }
  46
  47         $(
  48             pub(crate) struct $oty {
  49                 handle: handle::Handle,
  50                 // Prevent Send and Sync impls. `!Send`/`!Sync` is the usual
  51                 // way of doing this, but that requires unstable features.
  52                 // rust-analyzer uses this code and avoids unstable features.
  53                 _marker: PhantomData<*mut ()>,
  54             }
  55
  56             // Forward `Drop::drop` to the inherent `drop` method.
  57             impl Drop for $oty {
  58                 fn drop(&mut self) {
  59                     $oty {
  60                         handle: self.handle,
  61                         _marker: PhantomData,
  62                     }.drop();
  63                 }
  64             }
  65
  66             impl<S> Encode<S> for $oty {
  67                 fn encode(self, w: &mut Writer, s: &mut S) {
  68                     let handle = self.handle;
  69                     mem::forget(self);
  70                     handle.encode(w, s);
  71                 }
  72             }
  73
  74             impl<S: server::Types> DecodeMut<'_, '_, HandleStore<server::MarkedTypes<S>>>
  75                 for Marked<S::$oty, $oty>
  76             {
  77                 fn decode(r: &mut Reader<'_>, s: &mut HandleStore<server::MarkedTypes<S>>) -> Self {
  78                     s.$oty.take(handle::Handle::decode(r, &mut ()))
  79                 }
  80             }
  81
  82             impl<S> Encode<S> for &$oty {
  83                 fn encode(self, w: &mut Writer, s: &mut S) {
  84                     self.handle.encode(w, s);
  85                 }
  86             }
  87
  88             impl<'s, S: server::Types> Decode<'_, 's, HandleStore<server::MarkedTypes<S>>>
  89                 for &'s Marked<S::$oty, $oty>
  90             {
  91                 fn decode(r: &mut Reader<'_>, s: &'s HandleStore<server::MarkedTypes<S>>) -> Self {
  92                     &s.$oty[handle::Handle::decode(r, &mut ())]
  93                 }
  94             }
  95
  96             impl<S> Encode<S> for &mut $oty {
  97                 fn encode(self, w: &mut Writer, s: &mut S) {
  98                     self.handle.encode(w, s);
  99                 }
 100             }
 101
 102             impl<'s, S: server::Types> DecodeMut<'_, 's, HandleStore<server::MarkedTypes<S>>>
 103                 for &'s mut Marked<S::$oty, $oty>
 104             {
 105                 fn decode(
 106                     r: &mut Reader<'_>,
 107                     s: &'s mut HandleStore<server::MarkedTypes<S>>
 108                 ) -> Self {
 109                     &mut s.$oty[handle::Handle::decode(r, &mut ())]
 110                 }
 111             }
 112
 113             impl<S: server::Types> Encode<HandleStore<server::MarkedTypes<S>>>
 114                 for Marked<S::$oty, $oty>
 115             {
 116                 fn encode(self, w: &mut Writer, s: &mut HandleStore<server::MarkedTypes<S>>) {
 117                     s.$oty.alloc(self).encode(w, s);
 118                 }
 119             }
 120
 121             impl<S> DecodeMut<'_, '_, S> for $oty {
 122                 fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
 123                     $oty {
 124                         handle: handle::Handle::decode(r, s),
 125                         _marker: PhantomData,
 126                     }
 127                 }
 128             }
 129         )*
 130
 131         $(
 132             #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 133             pub(crate) struct $ity {
 134                 handle: handle::Handle,
 135                 // Prevent Send and Sync impls. `!Send`/`!Sync` is the usual
 136                 // way of doing this, but that requires unstable features.
 137                 // rust-analyzer uses this code and avoids unstable features.
 138                 _marker: PhantomData<*mut ()>,
 139             }
 140
 141             impl<S> Encode<S> for $ity {
 142                 fn encode(self, w: &mut Writer, s: &mut S) {
 143                     self.handle.encode(w, s);
 144                 }
 145             }
 146
 147             impl<S: server::Types> DecodeMut<'_, '_, HandleStore<server::MarkedTypes<S>>>
 148                 for Marked<S::$ity, $ity>
 149             {
 150                 fn decode(r: &mut Reader<'_>, s: &mut HandleStore<server::MarkedTypes<S>>) -> Self {
 151                     s.$ity.copy(handle::Handle::decode(r, &mut ()))
 152                 }
 153             }
 154
 155             impl<S: server::Types> Encode<HandleStore<server::MarkedTypes<S>>>
 156                 for Marked<S::$ity, $ity>
 157             {
 158                 fn encode(self, w: &mut Writer, s: &mut HandleStore<server::MarkedTypes<S>>) {
 159                     s.$ity.alloc(self).encode(w, s);
 160                 }
 161             }
 162
 163             impl<S> DecodeMut<'_, '_, S> for $ity {
 164                 fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
 165                     $ity {
 166                         handle: handle::Handle::decode(r, s),
 167                         _marker: PhantomData,
 168                     }
 169                 }
 170             }
 171         )*
 172     }
 173 }
 174 define_handles! {
 175     'owned:
 176     FreeFunctions,
 177     TokenStream,
 178     SourceFile,
 179     MultiSpan,
 180     Diagnostic,
 181
 182     'interned:
 183     Span,
 184 }
 185
 186 // FIXME(eddyb) generate these impls by pattern-matching on the
 187 // names of methods - also could use the presence of `fn drop`
 188 // to distinguish between 'owned and 'interned, above.
 189 // Alternatively, special "modes" could be listed of types in with_api
 190 // instead of pattern matching on methods, here and in server decl.
 191
 192 impl Clone for TokenStream {
 193     fn clone(&self) -> Self {
 194         self.clone()
 195     }
 196 }
 197
 198 impl Clone for SourceFile {
 199     fn clone(&self) -> Self {
 200         self.clone()
 201     }
 202 }
 203
 204 impl Span {
 205     pub(crate) fn def_site() -> Span {
 206         Bridge::with(|bridge| bridge.globals.def_site)
 207     }
 208
 209     pub(crate) fn call_site() -> Span {
 210         Bridge::with(|bridge| bridge.globals.call_site)
 211     }
 212
 213     pub(crate) fn mixed_site() -> Span {
 214         Bridge::with(|bridge| bridge.globals.mixed_site)
 215     }
 216 }
 217
 218 impl fmt::Debug for Span {
 219     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 220         f.write_str(&self.debug())
 221     }
 222 }
 223
 224 pub(crate) use super::symbol::Symbol;
 225
 226 macro_rules! define_client_side {
 227     ($($name:ident {
 228         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)*
 229     }),* $(,)?) => {
 230         $(impl $name {
 231             $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)* {
 232                 Bridge::with(|bridge| {
 233                     let mut buf = bridge.cached_buffer.take();
 234
 235                     buf.clear();
 236                     api_tags::Method::$name(api_tags::$name::$method).encode(&mut buf, &mut ());
 237                     reverse_encode!(buf; $($arg),*);
 238
 239                     buf = bridge.dispatch.call(buf);
 240
 241                     let r = Result::<_, PanicMessage>::decode(&mut &buf[..], &mut ());
 242
 243                     bridge.cached_buffer = buf;
 244
 245                     r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
 246                 })
 247             })*
 248         })*
 249     }
 250 }
 251 with_api!(self, self, define_client_side);
 252
 253 struct Bridge<'a> {
 254     /// Reusable buffer (only `clear`-ed, never shrunk), primarily
 255     /// used for making requests.
 256     cached_buffer: Buffer,
 257
 258     /// Server-side function that the client uses to make requests.
 259     dispatch: closure::Closure<'a, Buffer, Buffer>,
 260
 261     /// Provided globals for this macro expansion.
 262     globals: ExpnGlobals<Span>,
 263 }
 264
 265 impl<'a> !Send for Bridge<'a> {}
 266 impl<'a> !Sync for Bridge<'a> {}
 267
 268 enum BridgeState<'a> {
 269     /// No server is currently connected to this client.
 270     NotConnected,
 271
 272     /// A server is connected and available for requests.
 273     Connected(Bridge<'a>),
 274
 275     /// Access to the bridge is being exclusively acquired
 276     /// (e.g., during `BridgeState::with`).
 277     InUse,
 278 }
 279
 280 enum BridgeStateL {}
 281
 282 impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
 283     type Out = BridgeState<'a>;
 284 }
 285
 286 thread_local! {
 287     static BRIDGE_STATE: scoped_cell::ScopedCell<BridgeStateL> =
 288         scoped_cell::ScopedCell::new(BridgeState::NotConnected);
 289 }
 290
 291 impl BridgeState<'_> {
 292     /// Take exclusive control of the thread-local
 293     /// `BridgeState`, and pass it to `f`, mutably.
 294     /// The state will be restored after `f` exits, even
 295     /// by panic, including modifications made to it by `f`.
 296     ///
 297     /// N.B., while `f` is running, the thread-local state
 298     /// is `BridgeState::InUse`.
 299     fn with<R>(f: impl FnOnce(&mut BridgeState<'_>) -> R) -> R {
 300         BRIDGE_STATE.with(|state| {
 301             state.replace(BridgeState::InUse, |mut state| {
 302                 // FIXME(#52812) pass `f` directly to `replace` when `RefMutL` is gone
 303                 f(&mut *state)
 304             })
 305         })
 306     }
 307 }
 308
 309 impl Bridge<'_> {
 310     fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
 311         BridgeState::with(|state| match state {
 312             BridgeState::NotConnected => {
 313                 panic!("procedural macro API is used outside of a procedural macro");
 314             }
 315             BridgeState::InUse => {
 316                 panic!("procedural macro API is used while it's already in use");
 317             }
 318             BridgeState::Connected(bridge) => f(bridge),
 319         })
 320     }
 321 }
 322
 323 pub(crate) fn is_available() -> bool {
 324     BridgeState::with(|state| match state {
 325         BridgeState::Connected(_) | BridgeState::InUse => true,
 326         BridgeState::NotConnected => false,
 327     })
 328 }
 329
 330 /// A client-side RPC entry-point, which may be using a different `proc_macro`
 331 /// from the one used by the server, but can be invoked compatibly.
 332 ///
 333 /// Note that the (phantom) `I` ("input") and `O` ("output") type parameters
 334 /// decorate the `Client<I, O>` with the RPC "interface" of the entry-point, but
 335 /// do not themselves participate in ABI, at all, only facilitate type-checking.
 336 ///
 337 /// E.g. `Client<TokenStream, TokenStream>` is the common proc macro interface,
 338 /// used for `#[proc_macro] fn foo(input: TokenStream) -> TokenStream`,
 339 /// indicating that the RPC input and output will be serialized token streams,
 340 /// and forcing the use of APIs that take/return `S::TokenStream`, server-side.
 341 #[repr(C)]
 342 pub struct Client<I, O> {
 343     // FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
 344     // a wrapper `fn` pointer, once `const fn` can reference `static`s.
 345     pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
 346
 347     pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
 348
 349     pub(super) _marker: PhantomData<fn(I) -> O>,
 350 }
 351
 352 impl<I, O> Copy for Client<I, O> {}
 353 impl<I, O> Clone for Client<I, O> {
 354     fn clone(&self) -> Self {
 355         *self
 356     }
 357 }
 358
 359 fn maybe_install_panic_hook(force_show_panics: bool) {
 360     // Hide the default panic output within `proc_macro` expansions.
 361     // NB. the server can't do this because it may use a different libstd.
 362     static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
 363     HIDE_PANICS_DURING_EXPANSION.call_once(|| {
 364         let prev = panic::take_hook();
 365         panic::set_hook(Box::new(move |info| {
 366             let show = BridgeState::with(|state| match state {
 367                 BridgeState::NotConnected => true,
 368                 BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
 369             });
 370             if show {
 371                 prev(info)
 372             }
 373         }));
 374     });
 375 }
 376
 377 /// Client-side helper for handling client panics, entering the bridge,
 378 /// deserializing input and serializing output.
 379 // FIXME(eddyb) maybe replace `Bridge::enter` with this?
 380 fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
 381     config: BridgeConfig<'_>,
 382     f: impl FnOnce(A) -> R,
 383 ) -> Buffer {
 384     let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
 385
 386     panic::catch_unwind(panic::AssertUnwindSafe(|| {
 387         maybe_install_panic_hook(force_show_panics);
 388
 389         // Make sure the symbol store is empty before decoding inputs.
 390         Symbol::invalidate_all();
 391
 392         let reader = &mut &buf[..];
 393         let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
 394
 395         // Put the buffer we used for input back in the `Bridge` for requests.
 396         let new_state =
 397             BridgeState::Connected(Bridge { cached_buffer: buf.take(), dispatch, globals });
 398
 399         BRIDGE_STATE.with(|state| {
 400             state.set(new_state, || {
 401                 let output = f(input);
 402
 403                 // Take the `cached_buffer` back out, for the output value.
 404                 buf = Bridge::with(|bridge| bridge.cached_buffer.take());
 405
 406                 // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
 407                 // from encoding a panic (`Err(e: PanicMessage)`) to avoid
 408                 // having handles outside the `bridge.enter(|| ...)` scope, and
 409                 // to catch panics that could happen while encoding the success.
 410                 //
 411                 // Note that panics should be impossible beyond this point, but
 412                 // this is defensively trying to avoid any accidental panicking
 413                 // reaching the `extern "C"` (which should `abort` but might not
 414                 // at the moment, so this is also potentially preventing UB).
 415                 buf.clear();
 416                 Ok::<_, ()>(output).encode(&mut buf, &mut ());
 417             })
 418         })
 419     }))
 420     .map_err(PanicMessage::from)
 421     .unwrap_or_else(|e| {
 422         buf.clear();
 423         Err::<(), _>(e).encode(&mut buf, &mut ());
 424     });
 425
 426     // Now that a response has been serialized, invalidate all symbols
 427     // registered with the interner.
 428     Symbol::invalidate_all();
 429     buf
 430 }
 431
 432 impl Client<crate::TokenStream, crate::TokenStream> {
 433     pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
 434         Client {
 435             get_handle_counters: HandleCounters::get,
 436             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
 437                 run_client(bridge, |input| f(crate::TokenStream(Some(input))).0)
 438             }),
 439             _marker: PhantomData,
 440         }
 441     }
 442 }
 443
 444 impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
 445     pub const fn expand2(
 446         f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
 447     ) -> Self {
 448         Client {
 449             get_handle_counters: HandleCounters::get,
 450             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
 451                 run_client(bridge, |(input, input2)| {
 452                     f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
 453                 })
 454             }),
 455             _marker: PhantomData,
 456         }
 457     }
 458 }
 459
 460 #[repr(C)]
 461 #[derive(Copy, Clone)]
 462 pub enum ProcMacro {
 463     CustomDerive {
 464         trait_name: &'static str,
 465         attributes: &'static [&'static str],
 466         client: Client<crate::TokenStream, crate::TokenStream>,
 467     },
 468
 469     Attr {
 470         name: &'static str,
 471         client: Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream>,
 472     },
 473
 474     Bang {
 475         name: &'static str,
 476         client: Client<crate::TokenStream, crate::TokenStream>,
 477     },
 478 }
 479
 480 impl ProcMacro {
 481     pub fn name(&self) -> &'static str {
 482         match self {
 483             ProcMacro::CustomDerive { trait_name, .. } => trait_name,
 484             ProcMacro::Attr { name, .. } => name,
 485             ProcMacro::Bang { name, .. } => name,
 486         }
 487     }
 488
 489     pub const fn custom_derive(
 490         trait_name: &'static str,
 491         attributes: &'static [&'static str],
 492         expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
 493     ) -> Self {
 494         ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
 495     }
 496
 497     pub const fn attr(
 498         name: &'static str,
 499         expand: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
 500     ) -> Self {
 501         ProcMacro::Attr { name, client: Client::expand2(expand) }
 502     }
 503
 504     pub const fn bang(
 505         name: &'static str,
 506         expand: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy,
 507     ) -> Self {
 508         ProcMacro::Bang { name, client: Client::expand1(expand) }
 509     }
 510 }