Remove `TokenStreamBuilder`.

`TokenStreamBuilder` exists to concatenate multiple `TokenStream`s
together. This commit removes it, and moves the concatenation
functionality directly into `TokenStream`, via two new methods
`push_tree` and `push_stream`. This makes things both simpler and
faster.

`push_tree` is particularly important. `TokenStreamBuilder` only had a
single `push` method, which pushed a stream. But in practice most of the
time we push a single token tree rather than a stream, and `push_tree`
avoids the need to build a token stream with a single entry (which
requires two allocations, one for the `Lrc` and one for the `Vec`).

The main `push_tree` use arises from a change to one of the `ToInternal`
impls in `proc_macro_server.rs`. It now returns a `SmallVec` instead of
a `TokenStream`. This return value is then iterated over by
`concat_trees`, which does `push_tree` on each element. Furthermore, the
use of `SmallVec` avoids more allocations, because there is always only
one or two token trees.

Note: the removed `TokenStreamBuilder::push` method had some code to
deal with a quadratic blowup case from #57735. This commit removes the
code. I tried and failed to reproduce the blowup from that PR, before
and after this change. Various other changes have happened to
`TokenStreamBuilder` in the meantime, so I suspect the original problem
is no longer relevant, though I don't have proof of this. Generally
speaking, repeatedly extending a `Vec` without pre-determining its
capacity is *not* quadratic. It's also incredibly common, within rustc
and many other Rust programs, so if there were performance problems
there you'd think it would show up in other places, too.

1 files changed, 45 insertions, 72 deletions

diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
index 4d2049cbc41..1d3c4fcca0a 100644
--- a/compiler/rustc_ast/src/tokenstream.rs
+++ b/compiler/rustc_ast/src/tokenstream.rs
@@ -245,12 +245,12 @@ impl AttrTokenStream {
                                 // properly implemented - we always synthesize fake tokens,
                                 // so we never reach this code.
 
-                                let mut builder = TokenStreamBuilder::new();
+                                let mut stream = TokenStream::default();
                                 for inner_attr in inner_attrs {
-                                    builder.push(inner_attr.tokens());
+                                    stream.push_stream(inner_attr.tokens());
                                 }
-                                builder.push(delim_tokens.clone());
-                                *tree = TokenTree::Delimited(*span, *delim, builder.build());
+                                stream.push_stream(delim_tokens.clone());
+                                *tree = TokenTree::Delimited(*span, *delim, stream);
                                 found = true;
                                 break;
                             }
@@ -505,76 +505,49 @@ impl TokenStream {
 
         self.trees().map(|tree| TokenStream::flatten_token_tree(tree)).collect()
     }
-}
 
-// 99.5%+ of the time we have 1 or 2 elements in this vector.
-#[derive(Clone)]
-pub struct TokenStreamBuilder(SmallVec<[TokenStream; 2]>);
-
-impl TokenStreamBuilder {
-    pub fn new() -> TokenStreamBuilder {
-        TokenStreamBuilder(SmallVec::new())
-    }
-
-    pub fn push(&mut self, stream: TokenStream) {
-        self.0.push(stream);
-    }
-
-    pub fn build(self) -> TokenStream {
-        let mut streams = self.0;
-        match streams.len() {
-            0 => TokenStream::default(),
-            1 => streams.pop().unwrap(),
-            _ => {
-                // We will extend the first stream in `streams` with the
-                // elements from the subsequent streams. This requires using
-                // `make_mut()` on the first stream, and in practice this
-                // doesn't cause cloning 99.9% of the time.
-                //
-                // One very common use case is when `streams` has two elements,
-                // where the first stream has any number of elements within
-                // (often 1, but sometimes many more) and the second stream has
-                // a single element within.
-
-                // Determine how much the first stream will be extended.
-                // Needed to avoid quadratic blow up from on-the-fly
-                // reallocations (#57735).
-                let num_appends = streams.iter().skip(1).map(|ts| ts.len()).sum();
-
-                // Get the first stream, which will become the result stream.
-                // If it's `None`, create an empty stream.
-                let mut iter = streams.into_iter();
-                let mut res_stream_lrc = iter.next().unwrap().0;
-
-                // Append the subsequent elements to the result stream, after
-                // reserving space for them.
-                let res_vec_mut = Lrc::make_mut(&mut res_stream_lrc);
-                res_vec_mut.reserve(num_appends);
-                for stream in iter {
-                    let stream_iter = stream.0.iter().cloned();
-
-                    // If (a) `res_mut_vec` is not empty and the last tree
-                    // within it is a token tree marked with `Joint`, and (b)
-                    // `stream` is not empty and the first tree within it is a
-                    // token tree, and (c) the two tokens can be glued
-                    // together...
-                    if let Some(TokenTree::Token(last_tok, Spacing::Joint)) = res_vec_mut.last()
-                        && let Some(TokenTree::Token(tok, spacing)) = stream.0.first()
-                        && let Some(glued_tok) = last_tok.glue(&tok)
-                    {
-                        // ...then overwrite the last token tree in
-                        // `res_vec_mut` with the glued token, and skip the
-                        // first token tree from `stream`.
-                        *res_vec_mut.last_mut().unwrap() = TokenTree::Token(glued_tok, *spacing);
-                        res_vec_mut.extend(stream_iter.skip(1));
-                    } else {
-                        // Append all of `stream`.
-                        res_vec_mut.extend(stream_iter);
-                    }
-                }
+    // If `vec` is not empty, try to glue `tt` onto its last token. The return
+    // value indicates if gluing took place.
+    fn try_glue_to_last(vec: &mut Vec<TokenTree>, tt: &TokenTree) -> bool {
+        if let Some(TokenTree::Token(last_tok, Spacing::Joint)) = vec.last()
+            && let TokenTree::Token(tok, spacing) = tt
+            && let Some(glued_tok) = last_tok.glue(&tok)
+        {
+            // ...then overwrite the last token tree in `vec` with the
+            // glued token, and skip the first token tree from `stream`.
+            *vec.last_mut().unwrap() = TokenTree::Token(glued_tok, *spacing);
+            true
+        } else {
+            false
+        }
+    }
 
-                TokenStream(res_stream_lrc)
-            }
+    // Push `tt` onto the end of the stream, possibly gluing it to the last
+    // token. Uses `make_mut` to maximize efficiency.
+    pub fn push_tree(&mut self, tt: TokenTree) {
+        let vec_mut = Lrc::make_mut(&mut self.0);
+
+        if Self::try_glue_to_last(vec_mut, &tt) {
+            // nothing else to do
+        } else {
+            vec_mut.push(tt);
+        }
+    }
+
+    // Push `stream` onto the end of the stream, possibly gluing the first
+    // token tree to the last token. (No other token trees will be glued.)
+    // Uses `make_mut` to maximize efficiency.
+    pub fn push_stream(&mut self, stream: TokenStream) {
+        let vec_mut = Lrc::make_mut(&mut self.0);
+
+        let stream_iter = stream.0.iter().cloned();
+
+        if let Some(first) = stream.0.first() && Self::try_glue_to_last(vec_mut, first) {
+            // Now skip the first token tree from `stream`.
+            vec_mut.extend(stream_iter.skip(1));
+        } else {
+            // Append all of `stream`.
+            vec_mut.extend(stream_iter);
         }
     }
 }