Batch proc_macro RPC for TokenStream iteration and combination operations

compiler/rustc_expand/src/proc_macro_server.rs

@@ -277,12 +277,6 @@ impl ToInternal<rustc_errors::Level> for Level {
 
 pub struct FreeFunctions;
 
-#[derive(Clone)]
-pub struct TokenStreamIter {
-    cursor: tokenstream::Cursor,
-    stack: Vec<TokenTree<Group, Punct, Ident, Literal>>,
-}
-
 #[derive(Clone)]
 pub struct Group {
     delimiter: Delimiter,
@@ -382,8 +376,6 @@ impl<'a, 'b> Rustc<'a, 'b> {
 impl server::Types for Rustc<'_, '_> {
     type FreeFunctions = FreeFunctions;
     type TokenStream = TokenStream;
-    type TokenStreamBuilder = tokenstream::TokenStreamBuilder;
-    type TokenStreamIter = TokenStreamIter;
     type Group = Group;
     type Punct = Punct;
     type Ident = Ident;
@@ -408,9 +400,6 @@ impl server::FreeFunctions for Rustc<'_, '_> {
 }
 
 impl server::TokenStream for Rustc<'_, '_> {
-    fn new(&mut self) -> Self::TokenStream {
-        TokenStream::default()
-    }
     fn is_empty(&mut self, stream: &Self::TokenStream) -> bool {
         stream.is_empty()
     }
@@ -481,53 +470,74 @@ impl server::TokenStream for Rustc<'_, '_> {
     ) -> Self::TokenStream {
         tree.to_internal()
     }
-    fn into_iter(&mut self, stream: Self::TokenStream) -> Self::TokenStreamIter {
-        TokenStreamIter { cursor: stream.into_trees(), stack: vec![] }
-    }
-}
-
-impl server::TokenStreamBuilder for Rustc<'_, '_> {
-    fn new(&mut self) -> Self::TokenStreamBuilder {
-        tokenstream::TokenStreamBuilder::new()
-    }
-    fn push(&mut self, builder: &mut Self::TokenStreamBuilder, stream: Self::TokenStream) {
-        builder.push(stream);
+    fn concat_trees(
+        &mut self,
+        base: Option<Self::TokenStream>,
+        trees: Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>>,
+    ) -> Self::TokenStream {
+        let mut builder = tokenstream::TokenStreamBuilder::new();
+        if let Some(base) = base {
+            builder.push(base);
+        }
+        for tree in trees {
+            builder.push(tree.to_internal());
+        }
+        builder.build()
     }
-    fn build(&mut self, builder: Self::TokenStreamBuilder) -> Self::TokenStream {
+    fn concat_streams(
+        &mut self,
+        base: Option<Self::TokenStream>,
+        streams: Vec<Self::TokenStream>,
+    ) -> Self::TokenStream {
+        let mut builder = tokenstream::TokenStreamBuilder::new();
+        if let Some(base) = base {
+            builder.push(base);
+        }
+        for stream in streams {
+            builder.push(stream);
+        }
         builder.build()
     }
-}
-
-impl server::TokenStreamIter for Rustc<'_, '_> {
-    fn next(
+    fn into_iter(
         &mut self,
-        iter: &mut Self::TokenStreamIter,
-    ) -> Option<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
+        stream: Self::TokenStream,
+    ) -> Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
+        // FIXME: This is a raw port of the previous approach, and can probably
+        // be optimized.
+        let mut cursor = stream.into_trees();
+        let mut stack = Vec::new();
+        let mut tts = Vec::new();
         loop {
-            let tree = iter.stack.pop().or_else(|| {
-                let next = iter.cursor.next_with_spacing()?;
-                Some(TokenTree::from_internal((next, &mut iter.stack, self)))
-            })?;
-            // A hack used to pass AST fragments to attribute and derive macros
-            // as a single nonterminal token instead of a token stream.
-            // Such token needs to be "unwrapped" and not represented as a delimited group.
-            // FIXME: It needs to be removed, but there are some compatibility issues (see #73345).
-            if let TokenTree::Group(ref group) = tree {
-                if group.flatten {
-                    iter.cursor.append(group.stream.clone());
-                    continue;
+            let next = stack.pop().or_else(|| {
+                let next = cursor.next_with_spacing()?;
+                Some(TokenTree::from_internal((next, &mut stack, self)))
+            });
+            match next {
+                Some(TokenTree::Group(group)) => {
+                    // A hack used to pass AST fragments to attribute and derive
+                    // macros as a single nonterminal token instead of a token
+                    // stream.  Such token needs to be "unwrapped" and not
+                    // represented as a delimited group.
+                    // FIXME: It needs to be removed, but there are some
+                    // compatibility issues (see #73345).
+                    if group.flatten {
+                        cursor.append(group.stream);
+                        continue;
+                    }
+                    tts.push(TokenTree::Group(group));
                 }
+                Some(tt) => tts.push(tt),
+                None => return tts,
             }
-            return Some(tree);
         }
     }
 }
 
 impl server::Group for Rustc<'_, '_> {
-    fn new(&mut self, delimiter: Delimiter, stream: Self::TokenStream) -> Self::Group {
+    fn new(&mut self, delimiter: Delimiter, stream: Option<Self::TokenStream>) -> Self::Group {
         Group {
             delimiter,
-            stream,
+            stream: stream.unwrap_or_default(),
             span: DelimSpan::from_single(server::Span::call_site(self)),
             flatten: false,
         }

library/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,11 @@ 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
+                        .unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
+                })
             }),
             _marker: PhantomData,
         }
@@ -450,7 +446,9 @@ 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
+                        .unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
                 })
             }),
             _marker: PhantomData,

library/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>,
+                    trees: Vec<$S::TokenStream>,
+                ) -> $S::TokenStream;
+                fn into_iter(
+                    $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;
@@ -337,6 +333,21 @@ impl<T: Unmark, E: Unmark> Unmark for Result<T, E> {
     }
 }
 
+impl<T: Mark> Mark for Vec<T> {
+    type Unmarked = Vec<T::Unmarked>;
+    fn mark(unmarked: Self::Unmarked) -> Self {
+        // Should be a no-op due to std's in-place collect optimizations.
+        unmarked.into_iter().map(T::mark).collect()
+    }
+}
+impl<T: Unmark> Unmark for Vec<T> {
+    type Unmarked = Vec<T::Unmarked>;
+    fn unmark(self) -> Self::Unmarked {
+        // Should be a no-op due to std's in-place collect optimizations.
+        self.into_iter().map(T::unmark).collect()
+    }
+}
+
 macro_rules! mark_noop {
     ($($ty:ty),* $(,)?) => {
         $(
@@ -394,6 +405,39 @@ 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)*);
+    };
+}
+
 #[derive(Clone)]
 pub enum TokenTree<G, P, I, L> {
     Group(G),
@@ -402,30 +446,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),

library/proc_macro/src/bridge/rpc.rs

@@ -43,15 +43,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)),*
                 }
@@ -246,6 +248,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 {

library/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;