library/proc_macro/src/bridge/server.rs

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