library/proc_macro/src/bridge/client.rs

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