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