1
Fork 0

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.
This commit is contained in:
Nicholas Nethercote 2019-10-08 10:34:53 +11:00
parent 421bd77f42
commit 3832a634d3

View file

@ -249,20 +249,47 @@ impl TokenStream {
0 => TokenStream::empty(), 0 => TokenStream::empty(),
1 => streams.pop().unwrap(), 1 => streams.pop().unwrap(),
_ => { _ => {
// rust-lang/rust#57735: pre-allocate vector to avoid // We are going to extend the first stream in `streams` with
// quadratic blow-up due to on-the-fly reallocations. // the elements from the subsequent streams. This requires
let tree_count = streams.iter() // using `make_mut()` on the first stream, and in practice this
.map(|ts| match &ts.0 { None => 0, Some(s) => s.len() }) // doesn't cause cloning 99.9% of the time.
.sum(); //
let mut vec = Vec::with_capacity(tree_count); // 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.
for stream in streams { // Determine how much the first stream will be extended.
match stream.0 { // Needed to avoid quadratic blow up from on-the-fly
None => {}, // reallocations (#57735).
Some(stream2) => vec.extend(stream2.iter().cloned()), let num_appends = streams.iter()
.skip(1)
.map(|ts| ts.len())
.sum();
// 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)),
}
} }
} }
} }