Auto merge of #98186 - mystor:tokenstream_as_vec_tt, r=eddyb

bors · bors · commit 1d5f643300bf · 2022-06-18T07:37:14.000Z
Batch proc_macro RPC for TokenStream iteration and combination operations This is the first part of #86822, split off as requested in rust-lang/rust#86822 (review). It reduces the number of RPC calls required for common operations such as iterating over and concatenating TokenStreams.
diff --git a/proc_macro/src/bridge/client.rs b/proc_macro/src/bridge/client.rs
@@ -178,8 +178,6 @@ define_handles! {
     'owned:
     FreeFunctions,
     TokenStream,
-    TokenStreamBuilder,
-    TokenStreamIter,
     Group,
     Literal,
     SourceFile,
@@ -204,12 +202,6 @@ impl Clone for TokenStream {
     }
 }
 
-impl Clone for TokenStreamIter {
-    fn clone(&self) -> Self {
-        self.clone()
-    }
-}
-
 impl Clone for Group {
     fn clone(&self) -> Self {
         self.clone()
@@ -435,7 +427,7 @@ impl Client<crate::TokenStream, crate::TokenStream> {
         Client {
             get_handle_counters: HandleCounters::get,
             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
-                run_client(bridge, |input| f(crate::TokenStream(input)).0)
+                run_client(bridge, |input| f(crate::TokenStream(Some(input))).0)
             }),
             _marker: PhantomData,
         }
@@ -450,7 +442,7 @@ impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
             get_handle_counters: HandleCounters::get,
             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
                 run_client(bridge, |(input, input2)| {
-                    f(crate::TokenStream(input), crate::TokenStream(input2)).0
+                    f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
                 })
             }),
             _marker: PhantomData,
diff --git a/proc_macro/src/bridge/mod.rs b/proc_macro/src/bridge/mod.rs
@@ -60,33 +60,29 @@ macro_rules! with_api {
             TokenStream {
                 fn drop($self: $S::TokenStream);
                 fn clone($self: &$S::TokenStream) -> $S::TokenStream;
-                fn new() -> $S::TokenStream;
                 fn is_empty($self: &$S::TokenStream) -> bool;
                 fn expand_expr($self: &$S::TokenStream) -> Result<$S::TokenStream, ()>;
                 fn from_str(src: &str) -> $S::TokenStream;
                 fn to_string($self: &$S::TokenStream) -> String;
                 fn from_token_tree(
                     tree: TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>,
                 ) -> $S::TokenStream;
-                fn into_iter($self: $S::TokenStream) -> $S::TokenStreamIter;
-            },
-            TokenStreamBuilder {
-                fn drop($self: $S::TokenStreamBuilder);
-                fn new() -> $S::TokenStreamBuilder;
-                fn push($self: &mut $S::TokenStreamBuilder, stream: $S::TokenStream);
-                fn build($self: $S::TokenStreamBuilder) -> $S::TokenStream;
-            },
-            TokenStreamIter {
-                fn drop($self: $S::TokenStreamIter);
-                fn clone($self: &$S::TokenStreamIter) -> $S::TokenStreamIter;
-                fn next(
-                    $self: &mut $S::TokenStreamIter,
-                ) -> Option<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
+                fn concat_trees(
+                    base: Option<$S::TokenStream>,
+                    trees: Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>,
+                ) -> $S::TokenStream;
+                fn concat_streams(
+                    base: Option<$S::TokenStream>,
+                    streams: Vec<$S::TokenStream>,
+                ) -> $S::TokenStream;
+                fn into_trees(
+                    $self: $S::TokenStream
+                ) -> Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
             },
             Group {
                 fn drop($self: $S::Group);
                 fn clone($self: &$S::Group) -> $S::Group;
-                fn new(delimiter: Delimiter, stream: $S::TokenStream) -> $S::Group;
+                fn new(delimiter: Delimiter, stream: Option<$S::TokenStream>) -> $S::Group;
                 fn delimiter($self: &$S::Group) -> Delimiter;
                 fn stream($self: &$S::Group) -> $S::TokenStream;
                 fn span($self: &$S::Group) -> $S::Span;
@@ -311,29 +307,18 @@ impl<'a, T, M> Unmark for &'a mut Marked<T, M> {
     }
 }
 
-impl<T: Mark> Mark for Option<T> {
-    type Unmarked = Option<T::Unmarked>;
-    fn mark(unmarked: Self::Unmarked) -> Self {
-        unmarked.map(T::mark)
-    }
-}
-impl<T: Unmark> Unmark for Option<T> {
-    type Unmarked = Option<T::Unmarked>;
-    fn unmark(self) -> Self::Unmarked {
-        self.map(T::unmark)
-    }
-}
-
-impl<T: Mark, E: Mark> Mark for Result<T, E> {
-    type Unmarked = Result<T::Unmarked, E::Unmarked>;
+impl<T: Mark> Mark for Vec<T> {
+    type Unmarked = Vec<T::Unmarked>;
     fn mark(unmarked: Self::Unmarked) -> Self {
-        unmarked.map(T::mark).map_err(E::mark)
+        // Should be a no-op due to std's in-place collect optimizations.
+        unmarked.into_iter().map(T::mark).collect()
     }
 }
-impl<T: Unmark, E: Unmark> Unmark for Result<T, E> {
-    type Unmarked = Result<T::Unmarked, E::Unmarked>;
+impl<T: Unmark> Unmark for Vec<T> {
+    type Unmarked = Vec<T::Unmarked>;
     fn unmark(self) -> Self::Unmarked {
-        self.map(T::unmark).map_err(E::unmark)
+        // Should be a no-op due to std's in-place collect optimizations.
+        self.into_iter().map(T::unmark).collect()
     }
 }
 
@@ -367,7 +352,6 @@ mark_noop! {
     Level,
     LineColumn,
     Spacing,
-    Bound<usize>,
 }
 
 rpc_encode_decode!(
@@ -394,6 +378,61 @@ rpc_encode_decode!(
     }
 );
 
+macro_rules! mark_compound {
+    (enum $name:ident <$($T:ident),+> { $($variant:ident $(($field:ident))?),* $(,)? }) => {
+        impl<$($T: Mark),+> Mark for $name <$($T),+> {
+            type Unmarked = $name <$($T::Unmarked),+>;
+            fn mark(unmarked: Self::Unmarked) -> Self {
+                match unmarked {
+                    $($name::$variant $(($field))? => {
+                        $name::$variant $((Mark::mark($field)))?
+                    })*
+                }
+            }
+        }
+
+        impl<$($T: Unmark),+> Unmark for $name <$($T),+> {
+            type Unmarked = $name <$($T::Unmarked),+>;
+            fn unmark(self) -> Self::Unmarked {
+                match self {
+                    $($name::$variant $(($field))? => {
+                        $name::$variant $((Unmark::unmark($field)))?
+                    })*
+                }
+            }
+        }
+    }
+}
+
+macro_rules! compound_traits {
+    ($($t:tt)*) => {
+        rpc_encode_decode!($($t)*);
+        mark_compound!($($t)*);
+    };
+}
+
+compound_traits!(
+    enum Bound<T> {
+        Included(x),
+        Excluded(x),
+        Unbounded,
+    }
+);
+
+compound_traits!(
+    enum Option<T> {
+        Some(t),
+        None,
+    }
+);
+
+compound_traits!(
+    enum Result<T, E> {
+        Ok(t),
+        Err(e),
+    }
+);
+
 #[derive(Clone)]
 pub enum TokenTree<G, P, I, L> {
     Group(G),
@@ -402,30 +441,7 @@ pub enum TokenTree<G, P, I, L> {
     Literal(L),
 }
 
-impl<G: Mark, P: Mark, I: Mark, L: Mark> Mark for TokenTree<G, P, I, L> {
-    type Unmarked = TokenTree<G::Unmarked, P::Unmarked, I::Unmarked, L::Unmarked>;
-    fn mark(unmarked: Self::Unmarked) -> Self {
-        match unmarked {
-            TokenTree::Group(tt) => TokenTree::Group(G::mark(tt)),
-            TokenTree::Punct(tt) => TokenTree::Punct(P::mark(tt)),
-            TokenTree::Ident(tt) => TokenTree::Ident(I::mark(tt)),
-            TokenTree::Literal(tt) => TokenTree::Literal(L::mark(tt)),
-        }
-    }
-}
-impl<G: Unmark, P: Unmark, I: Unmark, L: Unmark> Unmark for TokenTree<G, P, I, L> {
-    type Unmarked = TokenTree<G::Unmarked, P::Unmarked, I::Unmarked, L::Unmarked>;
-    fn unmark(self) -> Self::Unmarked {
-        match self {
-            TokenTree::Group(tt) => TokenTree::Group(tt.unmark()),
-            TokenTree::Punct(tt) => TokenTree::Punct(tt.unmark()),
-            TokenTree::Ident(tt) => TokenTree::Ident(tt.unmark()),
-            TokenTree::Literal(tt) => TokenTree::Literal(tt.unmark()),
-        }
-    }
-}
-
-rpc_encode_decode!(
+compound_traits!(
     enum TokenTree<G, P, I, L> {
         Group(tt),
         Punct(tt),
diff --git a/proc_macro/src/bridge/rpc.rs b/proc_macro/src/bridge/rpc.rs
@@ -4,7 +4,6 @@ use std::any::Any;
 use std::char;
 use std::io::Write;
 use std::num::NonZeroU32;
-use std::ops::Bound;
 use std::str;
 
 pub(super) type Writer = super::buffer::Buffer;
@@ -43,15 +42,17 @@ macro_rules! rpc_encode_decode {
             }
         }
     };
-    (struct $name:ident { $($field:ident),* $(,)? }) => {
-        impl<S> Encode<S> for $name {
+    (struct $name:ident $(<$($T:ident),+>)? { $($field:ident),* $(,)? }) => {
+        impl<S, $($($T: Encode<S>),+)?> Encode<S> for $name $(<$($T),+>)? {
             fn encode(self, w: &mut Writer, s: &mut S) {
                 $(self.$field.encode(w, s);)*
             }
         }
 
-        impl<S> DecodeMut<'_, '_, S> for $name {
-            fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
+        impl<'a, S, $($($T: for<'s> DecodeMut<'a, 's, S>),+)?> DecodeMut<'a, '_, S>
+            for $name $(<$($T),+>)?
+        {
+            fn decode(r: &mut Reader<'a>, s: &mut S) -> Self {
                 $name {
                     $($field: DecodeMut::decode(r, s)),*
                 }
@@ -184,28 +185,6 @@ impl<'a, S, A: for<'s> DecodeMut<'a, 's, S>, B: for<'s> DecodeMut<'a, 's, S>> De
     }
 }
 
-rpc_encode_decode!(
-    enum Bound<T> {
-        Included(x),
-        Excluded(x),
-        Unbounded,
-    }
-);
-
-rpc_encode_decode!(
-    enum Option<T> {
-        None,
-        Some(x),
-    }
-);
-
-rpc_encode_decode!(
-    enum Result<T, E> {
-        Ok(x),
-        Err(e),
-    }
-);
-
 impl<S> Encode<S> for &[u8] {
     fn encode(self, w: &mut Writer, s: &mut S) {
         self.len().encode(w, s);
@@ -246,6 +225,26 @@ impl<S> DecodeMut<'_, '_, S> for String {
     }
 }
 
+impl<S, T: Encode<S>> Encode<S> for Vec<T> {
+    fn encode(self, w: &mut Writer, s: &mut S) {
+        self.len().encode(w, s);
+        for x in self {
+            x.encode(w, s);
+        }
+    }
+}
+
+impl<'a, S, T: for<'s> DecodeMut<'a, 's, S>> DecodeMut<'a, '_, S> for Vec<T> {
+    fn decode(r: &mut Reader<'a>, s: &mut S) -> Self {
+        let len = usize::decode(r, s);
+        let mut vec = Vec::with_capacity(len);
+        for _ in 0..len {
+            vec.push(T::decode(r, s));
+        }
+        vec
+    }
+}
+
 /// Simplified version of panic payloads, ignoring
 /// types other than `&'static str` and `String`.
 pub enum PanicMessage {
diff --git a/proc_macro/src/bridge/server.rs b/proc_macro/src/bridge/server.rs
@@ -8,8 +8,6 @@ use super::client::HandleStore;
 pub trait Types {
     type FreeFunctions: 'static;
     type TokenStream: 'static + Clone;
-    type TokenStreamBuilder: 'static;
-    type TokenStreamIter: 'static + Clone;
     type Group: 'static + Clone;
     type Punct: 'static + Copy + Eq + Hash;
     type Ident: 'static + Copy + Eq + Hash;
@@ -275,13 +273,17 @@ fn run_server<
 }
 
 impl client::Client<crate::TokenStream, crate::TokenStream> {
-    pub fn run<S: Server>(
+    pub fn run<S>(
         &self,
         strategy: &impl ExecutionStrategy,
         server: S,
         input: S::TokenStream,
         force_show_panics: bool,
-    ) -> Result<S::TokenStream, PanicMessage> {
+    ) -> Result<S::TokenStream, PanicMessage>
+    where
+        S: Server,
+        S::TokenStream: Default,
+    {
         let client::Client { get_handle_counters, run, _marker } = *self;
         run_server(
             strategy,
@@ -291,19 +293,23 @@ impl client::Client<crate::TokenStream, crate::TokenStream> {
             run,
             force_show_panics,
         )
-        .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
+        .map(|s| <Option<<MarkedTypes<S> as Types>::TokenStream>>::unmark(s).unwrap_or_default())
     }
 }
 
 impl client::Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
-    pub fn run<S: Server>(
+    pub fn run<S>(
         &self,
         strategy: &impl ExecutionStrategy,
         server: S,
         input: S::TokenStream,
         input2: S::TokenStream,
         force_show_panics: bool,
-    ) -> Result<S::TokenStream, PanicMessage> {
+    ) -> Result<S::TokenStream, PanicMessage>
+    where
+        S: Server,
+        S::TokenStream: Default,
+    {
         let client::Client { get_handle_counters, run, _marker } = *self;
         run_server(
             strategy,
@@ -316,6 +322,6 @@ impl client::Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream
             run,
             force_show_panics,
         )
-        .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
+        .map(|s| <Option<<MarkedTypes<S> as Types>::TokenStream>>::unmark(s).unwrap_or_default())
     }
 }
diff --git a/proc_macro/src/lib.rs b/proc_macro/src/lib.rs
