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         #[repr(C)]
  33         #[allow(non_snake_case)]
  34         pub(super) struct HandleStore<S: server::Types> {
  35             $($oty: handle::OwnedStore<S::$oty>,)*
  36             $($ity: handle::InternedStore<S::$ity>,)*
  37         }
  38
  39         impl<S: server::Types> HandleStore<S> {
  40             pub(super) fn new(handle_counters: &'static HandleCounters) -> Self {
  41                 HandleStore {
  42                     $($oty: handle::OwnedStore::new(&handle_counters.$oty),)*
  43                     $($ity: handle::InternedStore::new(&handle_counters.$ity),)*
  44                 }
  45             }
  46         }
  47
  48         $(
  49             #[repr(C)]
  50             pub(crate) struct $oty {
  51                 handle: handle::Handle,
  52                 // Prevent Send and Sync impls
  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             #[repr(C)]
 133             #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 134             pub(crate) struct $ity {
 135                 handle: handle::Handle,
 136                 // Prevent Send and Sync impls
 137                 _marker: PhantomData<*mut ()>,
 138             }
 139
 140             impl<S> Encode<S> for $ity {
 141                 fn encode(self, w: &mut Writer, s: &mut S) {
 142                     self.handle.encode(w, s);
 143                 }
 144             }
 145
 146             impl<S: server::Types> DecodeMut<'_, '_, HandleStore<server::MarkedTypes<S>>>
 147                 for Marked<S::$ity, $ity>
 148             {
 149                 fn decode(r: &mut Reader<'_>, s: &mut HandleStore<server::MarkedTypes<S>>) -> Self {
 150                     s.$ity.copy(handle::Handle::decode(r, &mut ()))
 151                 }
 152             }
 153
 154             impl<S: server::Types> Encode<HandleStore<server::MarkedTypes<S>>>
 155                 for Marked<S::$ity, $ity>
 156             {
 157                 fn encode(self, w: &mut Writer, s: &mut HandleStore<server::MarkedTypes<S>>) {
 158                     s.$ity.alloc(self).encode(w, s);
 159                 }
 160             }
 161
 162             impl<S> DecodeMut<'_, '_, S> for $ity {
 163                 fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
 164                     $ity {
 165                         handle: handle::Handle::decode(r, s),
 166                         _marker: PhantomData,
 167                     }
 168                 }
 169             }
 170         )*
 171     }
 172 }
 173 define_handles! {
 174     'owned:
 175     FreeFunctions,
 176     TokenStream,
 177     TokenStreamBuilder,
 178     TokenStreamIter,
 179     Group,
 180     Literal,
 181     SourceFile,
 182     MultiSpan,
 183     Diagnostic,
 184
 185     'interned:
 186     Punct,
 187     Ident,
 188     Span,
 189 }
 190
 191 // FIXME(eddyb) generate these impls by pattern-matching on the
 192 // names of methods - also could use the presence of `fn drop`
 193 // to distinguish between 'owned and 'interned, above.
 194 // Alternatively, special 'modes" could be listed of types in with_api
 195 // instead of pattern matching on methods, here and in server decl.
 196
 197 impl Clone for TokenStream {
 198     fn clone(&self) -> Self {
 199         self.clone()
 200     }
 201 }
 202
 203 impl Clone for TokenStreamIter {
 204     fn clone(&self) -> Self {
 205         self.clone()
 206     }
 207 }
 208
 209 impl Clone for Group {
 210     fn clone(&self) -> Self {
 211         self.clone()
 212     }
 213 }
 214
 215 impl Clone for Literal {
 216     fn clone(&self) -> Self {
 217         self.clone()
 218     }
 219 }
 220
 221 impl fmt::Debug for Literal {
 222     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 223         f.debug_struct("Literal")
 224             // format the kind without quotes, as in `kind: Float`
 225             .field("kind", &format_args!("{}", &self.debug_kind()))
 226             .field("symbol", &self.symbol())
 227             // format `Some("...")` on one line even in {:#?} mode
 228             .field("suffix", &format_args!("{:?}", &self.suffix()))
 229             .field("span", &self.span())
 230             .finish()
 231     }
 232 }
 233
 234 impl Clone for SourceFile {
 235     fn clone(&self) -> Self {
 236         self.clone()
 237     }
 238 }
 239
 240 impl fmt::Debug for Span {
 241     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 242         f.write_str(&self.debug())
 243     }
 244 }
 245
 246 macro_rules! define_client_side {
 247     ($($name:ident {
 248         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)*
 249     }),* $(,)?) => {
 250         $(impl $name {
 251             $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)* {
 252                 Bridge::with(|bridge| {
 253                     let mut b = bridge.cached_buffer.take();
 254
 255                     b.clear();
 256                     api_tags::Method::$name(api_tags::$name::$method).encode(&mut b, &mut ());
 257                     reverse_encode!(b; $($arg),*);
 258
 259                     b = bridge.dispatch.call(b);
 260
 261                     let r = Result::<_, PanicMessage>::decode(&mut &b[..], &mut ());
 262
 263                     bridge.cached_buffer = b;
 264
 265                     r.unwrap_or_else(|e| panic::resume_unwind(e.into()))
 266                 })
 267             })*
 268         })*
 269     }
 270 }
 271 with_api!(self, self, define_client_side);
 272
 273 enum BridgeState<'a> {
 274     /// No server is currently connected to this client.
 275     NotConnected,
 276
 277     /// A server is connected and available for requests.
 278     Connected(Bridge<'a>),
 279
 280     /// Access to the bridge is being exclusively acquired
 281     /// (e.g., during `BridgeState::with`).
 282     InUse,
 283 }
 284
 285 enum BridgeStateL {}
 286
 287 impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
 288     type Out = BridgeState<'a>;
 289 }
 290
 291 thread_local! {
 292     static BRIDGE_STATE: scoped_cell::ScopedCell<BridgeStateL> =
 293         scoped_cell::ScopedCell::new(BridgeState::NotConnected);
 294 }
 295
 296 impl BridgeState<'_> {
 297     /// Take exclusive control of the thread-local
 298     /// `BridgeState`, and pass it to `f`, mutably.
 299     /// The state will be restored after `f` exits, even
 300     /// by panic, including modifications made to it by `f`.
 301     ///
 302     /// N.B., while `f` is running, the thread-local state
 303     /// is `BridgeState::InUse`.
 304     fn with<R>(f: impl FnOnce(&mut BridgeState<'_>) -> R) -> R {
 305         BRIDGE_STATE.with(|state| {
 306             state.replace(BridgeState::InUse, |mut state| {
 307                 // FIXME(#52812) pass `f` directly to `replace` when `RefMutL` is gone
 308                 f(&mut *state)
 309             })
 310         })
 311     }
 312 }
 313
 314 impl Bridge<'_> {
 315     pub(crate) fn is_available() -> bool {
 316         BridgeState::with(|state| match state {
 317             BridgeState::Connected(_) | BridgeState::InUse => true,
 318             BridgeState::NotConnected => false,
 319         })
 320     }
 321
 322     fn enter<R>(self, f: impl FnOnce() -> R) -> R {
 323         let force_show_panics = self.force_show_panics;
 324         // Hide the default panic output within `proc_macro` expansions.
 325         // NB. the server can't do this because it may use a different libstd.
 326         static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
 327         HIDE_PANICS_DURING_EXPANSION.call_once(|| {
 328             let prev = panic::take_hook();
 329             panic::set_hook(Box::new(move |info| {
 330                 let show = BridgeState::with(|state| match state {
 331                     BridgeState::NotConnected => true,
 332                     BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
 333                 });
 334                 if show {
 335                     prev(info)
 336                 }
 337             }));
 338         });
 339
 340         BRIDGE_STATE.with(|state| state.set(BridgeState::Connected(self), f))
 341     }
 342
 343     fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
 344         BridgeState::with(|state| match state {
 345             BridgeState::NotConnected => {
 346                 panic!("procedural macro API is used outside of a procedural macro");
 347             }
 348             BridgeState::InUse => {
 349                 panic!("procedural macro API is used while it's already in use");
 350             }
 351             BridgeState::Connected(bridge) => f(bridge),
 352         })
 353     }
 354 }
 355
 356 /// A client-side "global object" (usually a function pointer),
 357 /// which may be using a different `proc_macro` from the one
 358 /// used by the server, but can be interacted with compatibly.
 359 ///
 360 /// N.B., `F` must have FFI-friendly memory layout (e.g., a pointer).
 361 /// The call ABI of function pointers used for `F` doesn't
 362 /// need to match between server and client, since it's only
 363 /// passed between them and (eventually) called by the client.
 364 #[repr(C)]
 365 #[derive(Copy, Clone)]
 366 pub struct Client<F> {
 367     // FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
 368     // a wrapper `fn` pointer, once `const fn` can reference `static`s.
 369     pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
 370     pub(super) run: extern "C" fn(Bridge<'_>, F) -> Buffer<u8>,
 371     pub(super) f: F,
 372 }
 373
 374 /// Client-side helper for handling client panics, entering the bridge,
 375 /// deserializing input and serializing output.
 376 // FIXME(eddyb) maybe replace `Bridge::enter` with this?
 377 fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
 378     mut bridge: Bridge<'_>,
 379     f: impl FnOnce(A) -> R,
 380 ) -> Buffer<u8> {
 381     // The initial `cached_buffer` contains the input.
 382     let mut b = bridge.cached_buffer.take();
 383
 384     panic::catch_unwind(panic::AssertUnwindSafe(|| {
 385         bridge.enter(|| {
 386             let reader = &mut &b[..];
 387             let input = A::decode(reader, &mut ());
 388
 389             // Put the `cached_buffer` back in the `Bridge`, for requests.
 390             Bridge::with(|bridge| bridge.cached_buffer = b.take());
 391
 392             let output = f(input);
 393
 394             // Take the `cached_buffer` back out, for the output value.
 395             b = Bridge::with(|bridge| bridge.cached_buffer.take());
 396
 397             // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
 398             // from encoding a panic (`Err(e: PanicMessage)`) to avoid
 399             // having handles outside the `bridge.enter(|| ...)` scope, and
 400             // to catch panics that could happen while encoding the success.
 401             //
 402             // Note that panics should be impossible beyond this point, but
 403             // this is defensively trying to avoid any accidental panicking
 404             // reaching the `extern "C"` (which should `abort` but might not
 405             // at the moment, so this is also potentially preventing UB).
 406             b.clear();
 407             Ok::<_, ()>(output).encode(&mut b, &mut ());
 408         })
 409     }))
 410     .map_err(PanicMessage::from)
 411     .unwrap_or_else(|e| {
 412         b.clear();
 413         Err::<(), _>(e).encode(&mut b, &mut ());
 414     });
 415     b
 416 }
 417
 418 impl Client<fn(crate::TokenStream) -> crate::TokenStream> {
 419     #[rustc_allow_const_fn_unstable(const_fn)]
 420     pub const fn expand1(f: fn(crate::TokenStream) -> crate::TokenStream) -> Self {
 421         extern "C" fn run(
 422             bridge: Bridge<'_>,
 423             f: impl FnOnce(crate::TokenStream) -> crate::TokenStream,
 424         ) -> Buffer<u8> {
 425             run_client(bridge, |input| f(crate::TokenStream(input)).0)
 426         }
 427         Client { get_handle_counters: HandleCounters::get, run, f }
 428     }
 429 }
 430
 431 impl Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream> {
 432     #[rustc_allow_const_fn_unstable(const_fn)]
 433     pub const fn expand2(
 434         f: fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
 435     ) -> Self {
 436         extern "C" fn run(
 437             bridge: Bridge<'_>,
 438             f: impl FnOnce(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
 439         ) -> Buffer<u8> {
 440             run_client(bridge, |(input, input2)| {
 441                 f(crate::TokenStream(input), crate::TokenStream(input2)).0
 442             })
 443         }
 444         Client { get_handle_counters: HandleCounters::get, run, f }
 445     }
 446 }
 447
 448 #[repr(C)]
 449 #[derive(Copy, Clone)]
 450 pub enum ProcMacro {
 451     CustomDerive {
 452         trait_name: &'static str,
 453         attributes: &'static [&'static str],
 454         client: Client<fn(crate::TokenStream) -> crate::TokenStream>,
 455     },
 456
 457     Attr {
 458         name: &'static str,
 459         client: Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream>,
 460     },
 461
 462     Bang {
 463         name: &'static str,
 464         client: Client<fn(crate::TokenStream) -> crate::TokenStream>,
 465     },
 466 }
 467
 468 impl ProcMacro {
 469     pub fn name(&self) -> &'static str {
 470         match self {
 471             ProcMacro::CustomDerive { trait_name, .. } => trait_name,
 472             ProcMacro::Attr { name, .. } => name,
 473             ProcMacro::Bang { name, .. } => name,
 474         }
 475     }
 476
 477     #[rustc_allow_const_fn_unstable(const_fn)]
 478     pub const fn custom_derive(
 479         trait_name: &'static str,
 480         attributes: &'static [&'static str],
 481         expand: fn(crate::TokenStream) -> crate::TokenStream,
 482     ) -> Self {
 483         ProcMacro::CustomDerive { trait_name, attributes, client: Client::expand1(expand) }
 484     }
 485
 486     #[rustc_allow_const_fn_unstable(const_fn)]
 487     pub const fn attr(
 488         name: &'static str,
 489         expand: fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
 490     ) -> Self {
 491         ProcMacro::Attr { name, client: Client::expand2(expand) }
 492     }
 493
 494     #[rustc_allow_const_fn_unstable(const_fn)]
 495     pub const fn bang(
 496         name: &'static str,
 497         expand: fn(crate::TokenStream) -> crate::TokenStream,
 498     ) -> Self {
 499         ProcMacro::Bang { name, client: Client::expand1(expand) }
 500     }
 501 }