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