proc_macro: reduce the number of messages required to create, extend, and iterate TokenStreams

mystor · mystor · commit 679335a962bb · 2022-06-17T00:42:26.000-04:00
This significantly reduces the cost of common interactions with TokenStream
when running with the CrossThread execution strategy, by reducing the number of
RPC calls required.
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,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,
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>,
+                    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;
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;
diff --git a/proc_macro/src/lib.rs b/proc_macro/src/lib.rs
@@ -43,7 +43,7 @@ use std::cmp::Ordering;
 use std::ops::RangeBounds;
 use std::path::PathBuf;
 use std::str::FromStr;
-use std::{error, fmt, iter, mem};
+use std::{error, fmt, iter};
 
 /// Determines whether proc_macro has been made accessible to the currently
 /// running program.
@@ -72,7 +72,7 @@ pub fn is_available() -> bool {
 /// and `#[proc_macro_derive]` definitions.
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 #[derive(Clone)]
-pub struct TokenStream(bridge::client::TokenStream);
+pub struct TokenStream(Option<bridge::client::TokenStream>);
 
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl !Send for TokenStream {}
@@ -126,13 +126,13 @@ impl TokenStream {
     /// Returns an empty `TokenStream` containing no token trees.
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
     pub fn new() -> TokenStream {
-        TokenStream(bridge::client::TokenStream::new())
+        TokenStream(None)
     }
 
     /// Checks if this `TokenStream` is empty.
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
     pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
+        self.0.as_ref().map(|h| h.is_empty()).unwrap_or(true)
     }
 
     /// Parses this `TokenStream` as an expression and attempts to expand any
@@ -147,8 +147,9 @@ impl TokenStream {
     /// considered errors, is unspecified and may change in the future.
     #[unstable(feature = "proc_macro_expand", issue = "90765")]
     pub fn expand_expr(&self) -> Result<TokenStream, ExpandError> {
-        match bridge::client::TokenStream::expand_expr(&self.0) {
-            Ok(stream) => Ok(TokenStream(stream)),
+        let stream = self.0.as_ref().ok_or(ExpandError)?;
+        match bridge::client::TokenStream::expand_expr(stream) {
+            Ok(stream) => Ok(TokenStream(Some(stream))),
             Err(_) => Err(ExpandError),
         }
     }
@@ -166,7 +167,7 @@ impl FromStr for TokenStream {
     type Err = LexError;
 
     fn from_str(src: &str) -> Result<TokenStream, LexError> {
-        Ok(TokenStream(bridge::client::TokenStream::from_str(src)))
+        Ok(TokenStream(Some(bridge::client::TokenStream::from_str(src))))
     }
 }
 
@@ -175,7 +176,7 @@ impl FromStr for TokenStream {
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl ToString for TokenStream {
     fn to_string(&self) -> String {
-        self.0.to_string()
+        self.0.as_ref().map(|t| t.to_string()).unwrap_or_default()
     }
 }
 
@@ -208,24 +209,38 @@ impl Default for TokenStream {
 #[unstable(feature = "proc_macro_quote", issue = "54722")]
 pub use quote::{quote, quote_span};
 
+fn tree_to_bridge_tree(
+    tree: TokenTree,
+) -> bridge::TokenTree<
+    bridge::client::Group,
+    bridge::client::Punct,
+    bridge::client::Ident,
+    bridge::client::Literal,
+> {
+    match tree {
+        TokenTree::Group(tt) => bridge::TokenTree::Group(tt.0),
+        TokenTree::Punct(tt) => bridge::TokenTree::Punct(tt.0),
+        TokenTree::Ident(tt) => bridge::TokenTree::Ident(tt.0),
+        TokenTree::Literal(tt) => bridge::TokenTree::Literal(tt.0),
+    }
+}
+
 /// Creates a token stream containing a single token tree.
 #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
 impl From<TokenTree> for TokenStream {
     fn from(tree: TokenTree) -> TokenStream {
-        TokenStream(bridge::client::TokenStream::from_token_tree(match tree {
-            TokenTree::Group(tt) => bridge::TokenTree::Group(tt.0),
-            TokenTree::Punct(tt) => bridge::TokenTree::Punct(tt.0),
-            TokenTree::Ident(tt) => bridge::TokenTree::Ident(tt.0),
-            TokenTree::Literal(tt) => bridge::TokenTree::Literal(tt.0),
-        }))
+        TokenStream(Some(bridge::client::TokenStream::from_token_tree(tree_to_bridge_tree(tree))))
     }
 }
 
 /// Collects a number of token trees into a single stream.
 #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
 impl iter::FromIterator<TokenTree> for TokenStream {
     fn from_iter<I: IntoIterator<Item = TokenTree>>(trees: I) -> Self {
-        trees.into_iter().map(TokenStream::from).collect()
+        TokenStream(Some(bridge::client::TokenStream::concat_trees(
+            None,
+            trees.into_iter().map(tree_to_bridge_tree).collect(),
+        )))
     }
 }
 
@@ -234,24 +249,30 @@ impl iter::FromIterator<TokenTree> for TokenStream {
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 impl iter::FromIterator<TokenStream> for TokenStream {
     fn from_iter<I: IntoIterator<Item = TokenStream>>(streams: I) -> Self {
-        let mut builder = bridge::client::TokenStreamBuilder::new();
-        streams.into_iter().for_each(|stream| builder.push(stream.0));
-        TokenStream(builder.build())
+        TokenStream(Some(bridge::client::TokenStream::concat_streams(
+            None,
+            streams.into_iter().filter_map(|stream| stream.0).collect(),
+        )))
     }
 }
 
 #[stable(feature = "token_stream_extend", since = "1.30.0")]
 impl Extend<TokenTree> for TokenStream {
     fn extend<I: IntoIterator<Item = TokenTree>>(&mut self, trees: I) {
-        self.extend(trees.into_iter().map(TokenStream::from));
+        *self = TokenStream(Some(bridge::client::TokenStream::concat_trees(
+            self.0.take(),
+            trees.into_iter().map(|tree| tree_to_bridge_tree(tree)).collect(),
+        )));
     }
 }
 
 #[stable(feature = "token_stream_extend", since = "1.30.0")]
 impl Extend<TokenStream> for TokenStream {
     fn extend<I: IntoIterator<Item = TokenStream>>(&mut self, streams: I) {
-        // FIXME(eddyb) Use an optimized implementation if/when possible.
-        *self = iter::once(mem::replace(self, Self::new())).chain(streams).collect();
+        *self = TokenStream(Some(bridge::client::TokenStream::concat_streams(
+            self.0.take(),
+            streams.into_iter().filter_map(|stream| stream.0).collect(),
+        )));
     }
 }
 
@@ -265,7 +286,16 @@ pub mod token_stream {
     /// and returns whole groups as token trees.
     #[derive(Clone)]
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
-    pub struct IntoIter(bridge::client::TokenStreamIter);
+    pub struct IntoIter(
+        std::vec::IntoIter<
+            bridge::TokenTree<
+                bridge::client::Group,
+                bridge::client::Punct,
+                bridge::client::Ident,
+                bridge::client::Literal,
+            >,
+        >,
+    );
 
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
     impl Iterator for IntoIter {
@@ -287,7 +317,7 @@ pub mod token_stream {
         type IntoIter = IntoIter;
 
         fn into_iter(self) -> IntoIter {
-            IntoIter(self.0.into_iter())
+            IntoIter(self.0.map(|v| v.into_iter()).unwrap_or_default().into_iter())
         }
     }
 }
@@ -734,7 +764,7 @@ impl Group {
     /// returned above.
     #[stable(feature = "proc_macro_lib2", since = "1.29.0")]
     pub fn stream(&self) -> TokenStream {
-        TokenStream(self.0.stream())
+        TokenStream(Some(self.0.stream()))
     }
 
     /// Returns the span for the delimiters of this token stream, spanning the
