src/libproc_macro/bridge/client.rs

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