Merge pull request #21154 from Veykril/push-nvlktkknrqms

proc-macro-srv: Fix unnecessary subtree wrapping in protocol

Author: Veykril

crates/proc-macro-api/src/legacy_protocol/msg/flat.rs

@@ -216,16 +216,7 @@ impl FlatTree {
             text: Vec::new(),
             version,
         };
-        let group = proc_macro_srv::Group {
-            delimiter: proc_macro_srv::Delimiter::None,
-            stream: Some(tokenstream),
-            span: proc_macro_srv::DelimSpan {
-                open: call_site,
-                close: call_site,
-                entire: call_site,
-            },
-        };
-        w.write_tokenstream(&group);
+        w.write_tokenstream(call_site, &tokenstream);
 
         FlatTree {
             subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
@@ -267,16 +258,7 @@ impl FlatTree {
             text: Vec::new(),
             version,
         };
-        let group = proc_macro_srv::Group {
-            delimiter: proc_macro_srv::Delimiter::None,
-            stream: Some(tokenstream),
-            span: proc_macro_srv::DelimSpan {
-                open: call_site,
-                close: call_site,
-                entire: call_site,
-            },
-        };
-        w.write_tokenstream(&group);
+        w.write_tokenstream(call_site, &tokenstream);
 
         FlatTree {
             subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
@@ -491,7 +473,7 @@ impl SpanTransformer for Span {
 }
 
 struct Writer<'a, 'span, S: SpanTransformer, W> {
-    work: VecDeque<(usize, W)>,
+    work: VecDeque<(usize, usize, W)>,
     string_table: FxHashMap<std::borrow::Cow<'a, str>, u32>,
     span_data_table: &'span mut S::Table,
     version: u32,
@@ -508,14 +490,13 @@ impl<'a, T: SpanTransformer> Writer<'a, '_, T, tt::iter::TtIter<'a, T::Span>> {
     fn write_subtree(&mut self, root: tt::SubtreeView<'a, T::Span>) {
         let subtree = root.top_subtree();
         self.enqueue(subtree, root.iter());
-        while let Some((idx, subtree)) = self.work.pop_front() {
-            self.subtree(idx, subtree);
+        while let Some((idx, len, subtree)) = self.work.pop_front() {
+            self.subtree(idx, len, subtree);
         }
     }
 
-    fn subtree(&mut self, idx: usize, subtree: tt::iter::TtIter<'a, T::Span>) {
+    fn subtree(&mut self, idx: usize, n_tt: usize, subtree: tt::iter::TtIter<'a, T::Span>) {
         let mut first_tt = self.token_tree.len();
-        let n_tt = subtree.clone().count(); // FIXME: `count()` walks over the entire iterator.
         self.token_tree.resize(first_tt + n_tt, !0);
 
         self.subtree[idx].tt = [first_tt as u32, (first_tt + n_tt) as u32];
@@ -594,7 +575,8 @@ impl<'a, T: SpanTransformer> Writer<'a, '_, T, tt::iter::TtIter<'a, T::Span>> {
         let close = self.token_id_of(subtree.delimiter.close);
         let delimiter_kind = subtree.delimiter.kind;
         self.subtree.push(SubtreeRepr { open, close, kind: delimiter_kind, tt: [!0, !0] });
-        self.work.push_back((idx, contents));
+        // FIXME: `count()` walks over the entire iterator.
+        self.work.push_back((idx, contents.clone().count(), contents));
         idx as u32
     }
 }
@@ -624,26 +606,43 @@ impl<'a, T: SpanTransformer, U> Writer<'a, '_, T, U> {
 }
 
 #[cfg(feature = "sysroot-abi")]
-impl<'a, T: SpanTransformer> Writer<'a, '_, T, &'a proc_macro_srv::Group<T::Span>> {
-    fn write_tokenstream(&mut self, root: &'a proc_macro_srv::Group<T::Span>) {
-        self.enqueue_group(root);
-
-        while let Some((idx, group)) = self.work.pop_front() {
-            self.group(idx, group);
+impl<'a, T: SpanTransformer>
+    Writer<'a, '_, T, Option<proc_macro_srv::TokenStreamIter<'a, T::Span>>>
+{
+    fn write_tokenstream(
+        &mut self,
+        call_site: T::Span,
+        root: &'a proc_macro_srv::TokenStream<T::Span>,
+    ) {
+        let call_site = self.token_id_of(call_site);
+        self.subtree.push(SubtreeRepr {
+            open: call_site,
+            close: call_site,
+            kind: tt::DelimiterKind::Invisible,
+            tt: [!0, !0],
+        });
+        self.work.push_back((0, root.len(), Some(root.iter())));
+
+        while let Some((idx, len, group)) = self.work.pop_front() {
+            self.group(idx, len, group);
         }
     }
 
-    fn group(&mut self, idx: usize, group: &'a proc_macro_srv::Group<T::Span>) {
+    fn group(
+        &mut self,
+        idx: usize,
+        n_tt: usize,
+        group: Option<proc_macro_srv::TokenStreamIter<'a, T::Span>>,
+    ) {
         let mut first_tt = self.token_tree.len();
-        let n_tt = group.stream.as_ref().map_or(0, |it| it.len());
         self.token_tree.resize(first_tt + n_tt, !0);
 
         self.subtree[idx].tt = [first_tt as u32, (first_tt + n_tt) as u32];
 
-        for tt in group.stream.iter().flat_map(|it| it.iter()) {
+        for tt in group.into_iter().flatten() {
             let idx_tag = match tt {
                 proc_macro_srv::TokenTree::Group(group) => {
-                    let idx = self.enqueue_group(group);
+                    let idx = self.enqueue(group);
                     idx << 2
                 }
                 proc_macro_srv::TokenTree::Literal(lit) => {
@@ -706,7 +705,7 @@ impl<'a, T: SpanTransformer> Writer<'a, '_, T, &'a proc_macro_srv::Group<T::Span
         }
     }
 
-    fn enqueue_group(&mut self, group: &'a proc_macro_srv::Group<T::Span>) -> u32 {
+    fn enqueue(&mut self, group: &'a proc_macro_srv::Group<T::Span>) -> u32 {
         let idx = self.subtree.len();
         let open = self.token_id_of(group.span.open);
         let close = self.token_id_of(group.span.close);
@@ -717,7 +716,11 @@ impl<'a, T: SpanTransformer> Writer<'a, '_, T, &'a proc_macro_srv::Group<T::Span
             proc_macro_srv::Delimiter::None => tt::DelimiterKind::Invisible,
         };
         self.subtree.push(SubtreeRepr { open, close, kind: delimiter_kind, tt: [!0, !0] });
-        self.work.push_back((idx, group));
+        self.work.push_back((
+            idx,
+            group.stream.as_ref().map_or(0, |stream| stream.len()),
+            group.stream.as_ref().map(|ts| ts.iter()),
+        ));
         idx as u32
     }
 }
@@ -959,7 +962,6 @@ impl<T: SpanTransformer> Reader<'_, T> {
             };
             res[i] = Some(g);
         }
-        // FIXME: double check this
         proc_macro_srv::TokenStream::new(vec![proc_macro_srv::TokenTree::Group(
             res[0].take().unwrap(),
         )])

crates/proc-macro-srv/src/lib.rs

@@ -56,7 +56,7 @@ pub use proc_macro::Delimiter;
 
 pub use crate::bridge::*;
 pub use crate::server_impl::literal_from_str;
-pub use crate::token_stream::{TokenStream, literal_to_string};
+pub use crate::token_stream::{TokenStream, TokenStreamIter, literal_to_string};
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 pub enum ProcMacroKind {