Skip to content

Commit 61ccbe1

Browse files
committed
Try to reduce codegen complexity of TokenStream's FromIterator and Extend impls
This is an experimental patch to try to reduce the codegen complexity of TokenStream's FromIterator and Extend implementations for downstream crates, by moving the core logic into a helper type. This might help improve build performance of crates which depend on proc_macro as iterators are used less, and the compiler may take less time to do things like attempt specializations or other iterator optimizations. The change intentionally sacrifices some optimization opportunities, such as using the specializations for collecting iterators derived from Vec::into_iter() into Vec. This is one of the simpler potential approaches to reducing the amount of code generated in crates depending on proc_macro, so it seems worth trying before other more-involved changes.
1 parent 0260d47 commit 61ccbe1

File tree

1 file changed

+91
-16
lines changed

1 file changed

+91
-16
lines changed

library/proc_macro/src/lib.rs

Lines changed: 91 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -230,14 +230,89 @@ impl From<TokenTree> for TokenStream {
230230
}
231231
}
232232

233+
/// Non-generic helper for implementing `FromIterator<TokenTree>` and
234+
/// `Extend<TokenTree>` with less monomorphization in calling crates.
235+
struct ExtendStreamWithTreesHelper {
236+
trees: Vec<
237+
bridge::TokenTree<
238+
bridge::client::TokenStream,
239+
bridge::client::Span,
240+
bridge::client::Symbol,
241+
>,
242+
>,
243+
}
244+
245+
impl ExtendStreamWithTreesHelper {
246+
fn new(capacity: usize) -> Self {
247+
ExtendStreamWithTreesHelper { trees: Vec::with_capacity(capacity) }
248+
}
249+
250+
fn push(&mut self, tree: TokenTree) {
251+
self.trees.push(tree_to_bridge_tree(tree));
252+
}
253+
254+
fn build(self) -> TokenStream {
255+
if self.trees.is_empty() {
256+
TokenStream(None)
257+
} else {
258+
TokenStream(Some(bridge::client::TokenStream::concat_trees(None, self.trees)))
259+
}
260+
}
261+
262+
fn extend(self, stream: &mut TokenStream) {
263+
if self.trees.is_empty() {
264+
return;
265+
}
266+
stream.0 = Some(bridge::client::TokenStream::concat_trees(stream.0.take(), self.trees))
267+
}
268+
}
269+
270+
/// Non-generic helper for implementing `FromIterator<TokenStream>` and
271+
/// `Extend<TokenStream>` with less monomorphization in calling crates.
272+
struct ExtendStreamWithStreamsHelper {
273+
streams: Vec<bridge::client::TokenStream>,
274+
}
275+
276+
impl ExtendStreamWithStreamsHelper {
277+
fn new(capacity: usize) -> Self {
278+
ExtendStreamWithStreamsHelper { streams: Vec::with_capacity(capacity) }
279+
}
280+
281+
fn push(&mut self, stream: TokenStream) {
282+
if let Some(stream) = stream.0 {
283+
self.streams.push(stream);
284+
}
285+
}
286+
287+
fn build(mut self) -> TokenStream {
288+
if self.streams.len() <= 1 {
289+
TokenStream(self.streams.pop())
290+
} else {
291+
TokenStream(Some(bridge::client::TokenStream::concat_streams(None, self.streams)))
292+
}
293+
}
294+
295+
fn extend(mut self, stream: &mut TokenStream) {
296+
if self.streams.is_empty() {
297+
return;
298+
}
299+
let base = stream.0.take();
300+
if base.is_none() && self.streams.len() == 1 {
301+
stream.0 = self.streams.pop();
302+
} else {
303+
stream.0 = Some(bridge::client::TokenStream::concat_streams(base, self.streams));
304+
}
305+
}
306+
}
307+
233308
/// Collects a number of token trees into a single stream.
234309
#[stable(feature = "proc_macro_lib2", since = "1.29.0")]
235310
impl iter::FromIterator<TokenTree> for TokenStream {
236311
fn from_iter<I: IntoIterator<Item = TokenTree>>(trees: I) -> Self {
237-
TokenStream(Some(bridge::client::TokenStream::concat_trees(
238-
None,
239-
trees.into_iter().map(tree_to_bridge_tree).collect(),
240-
)))
312+
let iter = trees.into_iter();
313+
let mut builder = ExtendStreamWithTreesHelper::new(iter.size_hint().0);
314+
iter.for_each(|tree| builder.push(tree));
315+
builder.build()
241316
}
242317
}
243318

@@ -246,30 +321,30 @@ impl iter::FromIterator<TokenTree> for TokenStream {
246321
#[stable(feature = "proc_macro_lib", since = "1.15.0")]
247322
impl iter::FromIterator<TokenStream> for TokenStream {
248323
fn from_iter<I: IntoIterator<Item = TokenStream>>(streams: I) -> Self {
249-
TokenStream(Some(bridge::client::TokenStream::concat_streams(
250-
None,
251-
streams.into_iter().filter_map(|stream| stream.0).collect(),
252-
)))
324+
let iter = streams.into_iter();
325+
let mut builder = ExtendStreamWithStreamsHelper::new(iter.size_hint().0);
326+
iter.for_each(|stream| builder.push(stream));
327+
builder.build()
253328
}
254329
}
255330

256331
#[stable(feature = "token_stream_extend", since = "1.30.0")]
257332
impl Extend<TokenTree> for TokenStream {
258333
fn extend<I: IntoIterator<Item = TokenTree>>(&mut self, trees: I) {
259-
*self = TokenStream(Some(bridge::client::TokenStream::concat_trees(
260-
self.0.take(),
261-
trees.into_iter().map(|tree| tree_to_bridge_tree(tree)).collect(),
262-
)));
334+
let iter = trees.into_iter();
335+
let mut builder = ExtendStreamWithTreesHelper::new(iter.size_hint().0);
336+
iter.for_each(|tree| builder.push(tree));
337+
builder.extend(self);
263338
}
264339
}
265340

266341
#[stable(feature = "token_stream_extend", since = "1.30.0")]
267342
impl Extend<TokenStream> for TokenStream {
268343
fn extend<I: IntoIterator<Item = TokenStream>>(&mut self, streams: I) {
269-
*self = TokenStream(Some(bridge::client::TokenStream::concat_streams(
270-
self.0.take(),
271-
streams.into_iter().filter_map(|stream| stream.0).collect(),
272-
)));
344+
let iter = streams.into_iter();
345+
let mut builder = ExtendStreamWithStreamsHelper::new(iter.size_hint().0);
346+
iter.for_each(|stream| builder.push(stream));
347+
builder.extend(self);
273348
}
274349
}
275350

0 commit comments

Comments
 (0)