crates/proc_macro_srv/src/proc_macro/bridge/server.rs

   1 //! lib-proc-macro server-side traits
   2 //!
   3 //! Copy from <https://github.com/rust-lang/rust/blob/6050e523bae6de61de4e060facc43dc512adaccd/src/libproc_macro/bridge/server.rs>
   4 //! augmented with removing unstable features
   5
   6 use super::*;
   7
   8 // FIXME(eddyb) generate the definition of `HandleStore` in `server.rs`.
   9 use super::client::HandleStore;
  10
  11 /// Declare an associated item of one of the traits below, optionally
  12 /// adjusting it (i.e., adding bounds to types and default bodies to methods).
  13 macro_rules! associated_item {
  14     (type FreeFunctions) =>
  15         (type FreeFunctions: 'static;);
  16     (type TokenStream) =>
  17         (type TokenStream: 'static + Clone;);
  18     (type TokenStreamBuilder) =>
  19         (type TokenStreamBuilder: 'static;);
  20     (type TokenStreamIter) =>
  21         (type TokenStreamIter: 'static + Clone;);
  22     (type Group) =>
  23         (type Group: 'static + Clone;);
  24     (type Punct) =>
  25         (type Punct: 'static + Copy + Eq + Hash;);
  26     (type Ident) =>
  27         (type Ident: 'static + Copy + Eq + Hash;);
  28     (type Literal) =>
  29         (type Literal: 'static + Clone;);
  30     (type SourceFile) =>
  31         (type SourceFile: 'static + Clone;);
  32     (type MultiSpan) =>
  33         (type MultiSpan: 'static;);
  34     (type Diagnostic) =>
  35         (type Diagnostic: 'static;);
  36     (type Span) =>
  37         (type Span: 'static + Copy + Eq + Hash;);
  38     (fn drop(&mut self, $arg:ident: $arg_ty:ty)) =>
  39         (fn drop(&mut self, $arg: $arg_ty) { mem::drop($arg) });
  40     (fn clone(&mut self, $arg:ident: $arg_ty:ty) -> $ret_ty:ty) =>
  41         (fn clone(&mut self, $arg: $arg_ty) -> $ret_ty { $arg.clone() });
  42     ($($item:tt)*) => ($($item)*;)
  43 }
  44
  45 macro_rules! declare_server_traits {
  46     ($($name:ident {
  47         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
  48     }),* $(,)?) => {
  49         pub trait Types {
  50             $(associated_item!(type $name);)*
  51         }
  52
  53         $(pub trait $name: Types {
  54             $(associated_item!(fn $method(&mut self, $($arg: $arg_ty),*) $(-> $ret_ty)?);)*
  55         })*
  56
  57         pub trait Server: Types $(+ $name)* {}
  58         impl<S: Types $(+ $name)*> Server for S {}
  59     }
  60 }
  61 with_api!(Self, self_, declare_server_traits);
  62
  63 pub(super) struct MarkedTypes<S: Types>(S);
  64
  65 macro_rules! define_mark_types_impls {
  66     ($($name:ident {
  67         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
  68     }),* $(,)?) => {
  69         impl<S: Types> Types for MarkedTypes<S> {
  70             $(type $name = Marked<S::$name, client::$name>;)*
  71         }
  72
  73         $(impl<S: $name> $name for MarkedTypes<S> {
  74             $(fn $method(&mut self, $($arg: $arg_ty),*) $(-> $ret_ty)? {
  75                 <_>::mark($name::$method(&mut self.0, $($arg.unmark()),*))
  76             })*
  77         })*
  78     }
  79 }
  80 with_api!(Self, self_, define_mark_types_impls);
  81
  82 struct Dispatcher<S: Types> {
  83     handle_store: HandleStore<S>,
  84     server: S,
  85 }
  86
  87 macro_rules! define_dispatcher_impl {
  88     ($($name:ident {
  89         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
  90     }),* $(,)?) => {
  91         // FIXME(eddyb) `pub` only for `ExecutionStrategy` below.
  92         pub trait DispatcherTrait {
  93             // HACK(eddyb) these are here to allow `Self::$name` to work below.
  94             $(type $name;)*
  95             fn dispatch(&mut self, b: Buffer<u8>) -> Buffer<u8>;
  96         }
  97
  98         impl<S: Server> DispatcherTrait for Dispatcher<MarkedTypes<S>> {
  99             $(type $name = <MarkedTypes<S> as Types>::$name;)*
 100             fn dispatch(&mut self, mut b: Buffer<u8>) -> Buffer<u8> {
 101                 let Dispatcher { handle_store, server } = self;
 102
 103                 let mut reader = &b[..];
 104                 match api_tags::Method::decode(&mut reader, &mut ()) {
 105                     $(api_tags::Method::$name(m) => match m {
 106                         $(api_tags::$name::$method => {
 107                             let mut call_method = || {
 108                                 reverse_decode!(reader, handle_store; $($arg: $arg_ty),*);
 109                                 $name::$method(server, $($arg),*)
 110                             };
 111                             // HACK(eddyb) don't use `panic::catch_unwind` in a panic.
 112                             // If client and server happen to use the same `libstd`,
 113                             // `catch_unwind` asserts that the panic counter was 0,
 114                             // even when the closure passed to it didn't panic.
 115                             let r = if thread::panicking() {
 116                                 Ok(call_method())
 117                             } else {
 118                                 panic::catch_unwind(panic::AssertUnwindSafe(call_method))
 119                                     .map_err(PanicMessage::from)
 120                             };
 121
 122                             b.clear();
 123                             r.encode(&mut b, handle_store);
 124                         })*
 125                     }),*
 126                 }
 127                 b
 128             }
 129         }
 130     }
 131 }
 132 with_api!(Self, self_, define_dispatcher_impl);
 133
 134 pub trait ExecutionStrategy {
 135     fn run_bridge_and_client<D: Copy + Send + 'static>(
 136         &self,
 137         dispatcher: &mut impl DispatcherTrait,
 138         input: Buffer<u8>,
 139         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
 140         client_data: D,
 141         force_show_panics: bool,
 142     ) -> Buffer<u8>;
 143 }
 144
 145 pub struct SameThread;
 146
 147 impl ExecutionStrategy for SameThread {
 148     fn run_bridge_and_client<D: Copy + Send + 'static>(
 149         &self,
 150         dispatcher: &mut impl DispatcherTrait,
 151         input: Buffer<u8>,
 152         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
 153         client_data: D,
 154         force_show_panics: bool,
 155     ) -> Buffer<u8> {
 156         let mut dispatch = |b| dispatcher.dispatch(b);
 157
 158         run_client(
 159             Bridge { cached_buffer: input, dispatch: (&mut dispatch).into(), force_show_panics },
 160             client_data,
 161         )
 162     }
 163 }
 164
 165 // NOTE(eddyb) Two implementations are provided, the second one is a bit
 166 // faster but neither is anywhere near as fast as same-thread execution.
 167
 168 pub struct CrossThread1;
 169
 170 impl ExecutionStrategy for CrossThread1 {
 171     fn run_bridge_and_client<D: Copy + Send + 'static>(
 172         &self,
 173         dispatcher: &mut impl DispatcherTrait,
 174         input: Buffer<u8>,
 175         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
 176         client_data: D,
 177         force_show_panics: bool,
 178     ) -> Buffer<u8> {
 179         use std::sync::mpsc::channel;
 180
 181         let (req_tx, req_rx) = channel();
 182         let (res_tx, res_rx) = channel();
 183
 184         let join_handle = thread::Builder::new()
 185             .name("Dispatch".to_owned())
 186             .spawn(move || {
 187                 let mut dispatch = |b| {
 188                     req_tx.send(b).unwrap();
 189                     res_rx.recv().unwrap()
 190                 };
 191
 192                 run_client(
 193                     Bridge {
 194                         cached_buffer: input,
 195                         dispatch: (&mut dispatch).into(),
 196                         force_show_panics,
 197                     },
 198                     client_data,
 199                 )
 200             })
 201             .expect("failed to spawn thread");
 202
 203         for b in req_rx {
 204             res_tx.send(dispatcher.dispatch(b)).unwrap();
 205         }
 206
 207         join_handle.join().unwrap()
 208     }
 209 }
 210
 211 pub struct CrossThread2;
 212
 213 impl ExecutionStrategy for CrossThread2 {
 214     fn run_bridge_and_client<D: Copy + Send + 'static>(
 215         &self,
 216         dispatcher: &mut impl DispatcherTrait,
 217         input: Buffer<u8>,
 218         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
 219         client_data: D,
 220         force_show_panics: bool,
 221     ) -> Buffer<u8> {
 222         use std::sync::{Arc, Mutex};
 223
 224         enum State<T> {
 225             Req(T),
 226             Res(T),
 227         }
 228
 229         let mut state = Arc::new(Mutex::new(State::Res(Buffer::new())));
 230
 231         let server_thread = thread::current();
 232         let state2 = state.clone();
 233         let join_handle = thread::Builder::new()
 234             .name("ProcMacroServer".to_owned())
 235             .spawn(move || {
 236                 let mut dispatch = |b| {
 237                     *state2.lock().unwrap() = State::Req(b);
 238                     server_thread.unpark();
 239                     loop {
 240                         thread::park();
 241                         if let State::Res(b) = &mut *state2.lock().unwrap() {
 242                             break b.take();
 243                         }
 244                     }
 245                 };
 246
 247                 let r = run_client(
 248                     Bridge {
 249                         cached_buffer: input,
 250                         dispatch: (&mut dispatch).into(),
 251                         force_show_panics,
 252                     },
 253                     client_data,
 254                 );
 255
 256                 // Wake up the server so it can exit the dispatch loop.
 257                 drop(state2);
 258                 server_thread.unpark();
 259
 260                 r
 261             })
 262             .expect("failed to spawn thread");
 263
 264         // Check whether `state2` was dropped, to know when to stop.
 265         while Arc::get_mut(&mut state).is_none() {
 266             thread::park();
 267             let mut b = match &mut *state.lock().unwrap() {
 268                 State::Req(b) => b.take(),
 269                 _ => continue,
 270             };
 271             b = dispatcher.dispatch(b.take());
 272             *state.lock().unwrap() = State::Res(b);
 273             join_handle.thread().unpark();
 274         }
 275
 276         join_handle.join().unwrap()
 277     }
 278 }
 279
 280 fn run_server<
 281     S: Server,
 282     I: Encode<HandleStore<MarkedTypes<S>>>,
 283     O: for<'a, 's> DecodeMut<'a, 's, HandleStore<MarkedTypes<S>>>,
 284     D: Copy + Send + 'static,
 285 >(
 286     strategy: &impl ExecutionStrategy,
 287     handle_counters: &'static client::HandleCounters,
 288     server: S,
 289     input: I,
 290     run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
 291     client_data: D,
 292     force_show_panics: bool,
 293 ) -> Result<O, PanicMessage> {
 294     let mut dispatcher =
 295         Dispatcher { handle_store: HandleStore::new(handle_counters), server: MarkedTypes(server) };
 296
 297     let mut b = Buffer::new();
 298     input.encode(&mut b, &mut dispatcher.handle_store);
 299
 300     b = strategy.run_bridge_and_client(
 301         &mut dispatcher,
 302         b,
 303         run_client,
 304         client_data,
 305         force_show_panics,
 306     );
 307
 308     Result::decode(&mut &b[..], &mut dispatcher.handle_store)
 309 }
 310
 311 impl client::Client<fn(crate::TokenStream) -> crate::TokenStream> {
 312     pub fn run<S: Server>(
 313         &self,
 314         strategy: &impl ExecutionStrategy,
 315         server: S,
 316         input: S::TokenStream,
 317         force_show_panics: bool,
 318     ) -> Result<S::TokenStream, PanicMessage> {
 319         let client::Client { get_handle_counters, run, f } = *self;
 320         run_server(
 321             strategy,
 322             get_handle_counters(),
 323             server,
 324             <MarkedTypes<S> as Types>::TokenStream::mark(input),
 325             run,
 326             f,
 327             force_show_panics,
 328         )
 329         .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
 330     }
 331 }
 332
 333 impl client::Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream> {
 334     pub fn run<S: Server>(
 335         &self,
 336         strategy: &impl ExecutionStrategy,
 337         server: S,
 338         input: S::TokenStream,
 339         input2: S::TokenStream,
 340         force_show_panics: bool,
 341     ) -> Result<S::TokenStream, PanicMessage> {
 342         let client::Client { get_handle_counters, run, f } = *self;
 343         run_server(
 344             strategy,
 345             get_handle_counters(),
 346             server,
 347             (
 348                 <MarkedTypes<S> as Types>::TokenStream::mark(input),
 349                 <MarkedTypes<S> as Types>::TokenStream::mark(input2),
 350             ),
 351             run,
 352             f,
 353             force_show_panics,
 354         )
 355         .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
 356     }
 357 }