crates/proc_macro_srv/src/proc_macro_nightly/bridge/client.rs

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