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