Optimize `TokenStream::from_streams`.

Currently, this function creates a new empty stream, and then appends
the elements from each given stream onto that stream. This can cause
quadratic behaviour.

This commit changes the function so that it modifies the first stream
(which can be long) by extending it with the elements from the
subsequent streams (which are almost always short), which avoids the
quadratic behaviour.

1 files changed, 37 insertions, 10 deletions

diff --git a/src/libsyntax/tokenstream.rs b/src/libsyntax/tokenstream.rs
index 26cae2a8e7c..76983f8ef85 100644
--- a/src/libsyntax/tokenstream.rs
+++ b/src/libsyntax/tokenstream.rs
@@ -249,20 +249,47 @@ impl TokenStream {
             0 => TokenStream::empty(),
             1 => streams.pop().unwrap(),
             _ => {
-                // rust-lang/rust#57735: pre-allocate vector to avoid
-                // quadratic blow-up due to on-the-fly reallocations.
-                let tree_count = streams.iter()
-                    .map(|ts| match &ts.0 { None => 0, Some(s) => s.len() })
+                // We are going to 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();
-                let mut vec = Vec::with_capacity(tree_count);
 
-                for stream in streams {
-                    match stream.0 {
-                        None => {},
-                        Some(stream2) => vec.extend(stream2.iter().cloned()),
+                // Get the first stream. If it's `None`, create an empty
+                // stream.
+                let mut iter = streams.drain();
+                let mut first_stream_lrc = match iter.next().unwrap().0 {
+                    Some(first_stream_lrc) => first_stream_lrc,
+                    None => Lrc::new(vec![]),
+                };
+
+                // Append the elements to the first stream, after reserving
+                // space for them.
+                let first_vec_mut = Lrc::make_mut(&mut first_stream_lrc);
+                first_vec_mut.reserve(num_appends);
+                for stream in iter {
+                    if let Some(stream) = stream.0 {
+                        first_vec_mut.extend(stream.iter().cloned());
                     }
                 }
-                TokenStream::new(vec)
+
+                // Create the final `TokenStream`.
+                match first_vec_mut.len() {
+                    0 => TokenStream(None),
+                    _ => TokenStream(Some(first_stream_lrc)),
+                }
             }
         }
     }