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