Skip to content

Commit 0a049fd

Browse files
mystormystor
committed
proc_macro: reduce the number of messages required to create, extend, and iterate TokenStreams
This significantly reduces the cost of common interactions with TokenStream when running with the CrossThread execution strategy, by reducing the number of RPC calls required.
1 parent 2b17219 commit 0a049fd

File tree

5 files changed

+128
-96
lines changed

5 files changed

+128
-96
lines changed

‎compiler/rustc_expand/src/proc_macro_server.rs‎

Lines changed: 54 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -277,12 +277,6 @@ impl ToInternal<rustc_errors::Level> for Level {
277277

278278
pub struct FreeFunctions;
279279

280-
#[derive(Clone)]
281-
pub struct TokenStreamIter {
282-
cursor: tokenstream::Cursor,
283-
stack: Vec<TokenTree<Group, Punct, Ident, Literal>>,
284-
}
285-
286280
#[derive(Clone)]
287281
pub struct Group {
288282
delimiter: Delimiter,
@@ -382,8 +376,6 @@ impl<'a, 'b> Rustc<'a, 'b> {
382376
impl server::Types for Rustc<'_, '_> {
383377
type FreeFunctions = FreeFunctions;
384378
type TokenStream = TokenStream;
385-
type TokenStreamBuilder = tokenstream::TokenStreamBuilder;
386-
type TokenStreamIter = TokenStreamIter;
387379
type Group = Group;
388380
type Punct = Punct;
389381
type Ident = Ident;
@@ -408,9 +400,6 @@ impl server::FreeFunctions for Rustc<'_, '_> {
408400
}
409401

410402
impl server::TokenStream for Rustc<'_, '_> {
411-
fn new(&mut self) -> Self::TokenStream {
412-
TokenStream::default()
413-
}
414403
fn is_empty(&mut self, stream: &Self::TokenStream) -> bool {
415404
stream.is_empty()
416405
}
@@ -481,53 +470,74 @@ impl server::TokenStream for Rustc<'_, '_> {
481470
) -> Self::TokenStream {
482471
tree.to_internal()
483472
}
484-
fn into_iter(&mut self, stream: Self::TokenStream) -> Self::TokenStreamIter {
485-
TokenStreamIter { cursor: stream.into_trees(), stack: vec![] }
486-
}
487-
}
488-
489-
impl server::TokenStreamBuilder for Rustc<'_, '_> {
490-
fn new(&mut self) -> Self::TokenStreamBuilder {
491-
tokenstream::TokenStreamBuilder::new()
492-
}
493-
fn push(&mut self, builder: &mut Self::TokenStreamBuilder, stream: Self::TokenStream) {
494-
builder.push(stream);
473+
fn concat_trees(
474+
&mut self,
475+
base: Option<Self::TokenStream>,
476+
trees: Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>>,
477+
) -> Self::TokenStream {
478+
let mut builder = tokenstream::TokenStreamBuilder::new();
479+
if let Some(base) = base {
480+
builder.push(base);
481+
}
482+
for tree in trees {
483+
builder.push(tree.to_internal());
484+
}
485+
builder.build()
495486
}
496-
fn build(&mut self, builder: Self::TokenStreamBuilder) -> Self::TokenStream {
487+
fn concat_streams(
488+
&mut self,
489+
base: Option<Self::TokenStream>,
490+
streams: Vec<Self::TokenStream>,
491+
) -> Self::TokenStream {
492+
let mut builder = tokenstream::TokenStreamBuilder::new();
493+
if let Some(base) = base {
494+
builder.push(base);
495+
}
496+
for stream in streams {
497+
builder.push(stream);
498+
}
497499
builder.build()
498500
}
499-
}
500-
501-
impl server::TokenStreamIter for Rustc<'_, '_> {
502-
fn next(
501+
fn into_iter(
503502
&mut self,
504-
iter: &mut Self::TokenStreamIter,
505-
) -> Option<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
503+
stream: Self::TokenStream,
504+
) -> Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
505+
// XXX: This is a raw port of the previous approach, and can probably be
506+
// optimized.
507+
let mut cursor = stream.into_trees();
508+
let mut stack = Vec::new();
509+
let mut tts = Vec::new();
506510
loop {
507-
let tree = iter.stack.pop().or_else(|| {
508-
let next = iter.cursor.next_with_spacing()?;
509-
Some(TokenTree::from_internal((next, &mut iter.stack, self)))
510-
})?;
511-
// A hack used to pass AST fragments to attribute and derive macros
512-
// as a single nonterminal token instead of a token stream.
513-
// Such token needs to be "unwrapped" and not represented as a delimited group.
514-
// FIXME: It needs to be removed, but there are some compatibility issues (see #73345).
515-
if let TokenTree::Group(ref group) = tree {
516-
if group.flatten {
517-
iter.cursor.append(group.stream.clone());
518-
continue;
511+
let next = stack.pop().or_else(|| {
512+
let next = cursor.next_with_spacing()?;
513+
Some(TokenTree::from_internal((next, &mut stack, self)))
514+
});
515+
match next {
516+
Some(TokenTree::Group(group)) => {
517+
// A hack used to pass AST fragments to attribute and derive
518+
// macros as a single nonterminal token instead of a token
519+
// stream. Such token needs to be "unwrapped" and not
520+
// represented as a delimited group.
521+
// FIXME: It needs to be removed, but there are some
522+
// compatibility issues (see #73345).
523+
if group.flatten {
524+
cursor.append(group.stream);
525+
continue;
526+
}
527+
tts.push(TokenTree::Group(group));
519528
}
529+
Some(tt) => tts.push(tt),
530+
None => return tts,
520531
}
521-
return Some(tree);
522532
}
523533
}
524534
}
525535

526536
impl server::Group for Rustc<'_, '_> {
527-
fn new(&mut self, delimiter: Delimiter, stream: Self::TokenStream) -> Self::Group {
537+
fn new(&mut self, delimiter: Delimiter, stream: Option<Self::TokenStream>) -> Self::Group {
528538
Group {
529539
delimiter,
530-
stream,
540+
stream: stream.unwrap_or_default(),
531541
span: DelimSpan::from_single(server::Span::call_site(self)),
532542
flatten: false,
533543
}

‎library/proc_macro/src/bridge/client.rs‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -178,8 +178,6 @@ define_handles! {
178178
'owned:
179179
FreeFunctions,
180180
TokenStream,
181-
TokenStreamBuilder,
182-
TokenStreamIter,
183181
Group,
184182
Literal,
185183
SourceFile,
@@ -204,12 +202,6 @@ impl Clone for TokenStream {
204202
}
205203
}
206204

207-
impl Clone for TokenStreamIter {
208-
fn clone(&self) -> Self {
209-
self.clone()
210-
}
211-
}
212-
213205
impl Clone for Group {
214206
fn clone(&self) -> Self {
215207
self.clone()
@@ -435,7 +427,11 @@ impl Client<crate::TokenStream, crate::TokenStream> {
435427
Client {
436428
get_handle_counters: HandleCounters::get,
437429
run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
438-
run_client(bridge, |input| f(crate::TokenStream(input)).0)
430+
run_client(bridge, |input| {
431+
f(crate::TokenStream(Some(input)))
432+
.0
433+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
434+
})
439435
}),
440436
_marker: PhantomData,
441437
}
@@ -450,7 +446,9 @@ impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
450446
get_handle_counters: HandleCounters::get,
451447
run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
452448
run_client(bridge, |(input, input2)| {
453-
f(crate::TokenStream(input), crate::TokenStream(input2)).0
449+
f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2)))
450+
.0
451+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
454452
})
455453
}),
456454
_marker: PhantomData,

‎library/proc_macro/src/bridge/mod.rs‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -60,33 +60,29 @@ macro_rules! with_api {
6060
TokenStream {
6161
fn drop($self: $S::TokenStream);
6262
fn clone($self: &$S::TokenStream) -> $S::TokenStream;
63-
fn new() -> $S::TokenStream;
6463
fn is_empty($self: &$S::TokenStream) -> bool;
6564
fn expand_expr($self: &$S::TokenStream) -> Result<$S::TokenStream, ()>;
6665
fn from_str(src: &str) -> $S::TokenStream;
6766
fn to_string($self: &$S::TokenStream) -> String;
6867
fn from_token_tree(
6968
tree: TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>,
7069
) -> $S::TokenStream;
71-
fn into_iter($self: $S::TokenStream) -> $S::TokenStreamIter;
72-
},
73-
TokenStreamBuilder {
74-
fn drop($self: $S::TokenStreamBuilder);
75-
fn new() -> $S::TokenStreamBuilder;
76-
fn push($self: &mut $S::TokenStreamBuilder, stream: $S::TokenStream);
77-
fn build($self: $S::TokenStreamBuilder) -> $S::TokenStream;
78-
},
79-
TokenStreamIter {
80-
fn drop($self: $S::TokenStreamIter);
81-
fn clone($self: &$S::TokenStreamIter) -> $S::TokenStreamIter;
82-
fn next(
83-
$self: &mut $S::TokenStreamIter,
84-
) -> Option<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
70+
fn concat_trees(
71+
base: Option<$S::TokenStream>,
72+
trees: Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>,
73+
) -> $S::TokenStream;
74+
fn concat_streams(
75+
base: Option<$S::TokenStream>,
76+
trees: Vec<$S::TokenStream>,
77+
) -> $S::TokenStream;
78+
fn into_iter(
79+
$self: $S::TokenStream
80+
) -> Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
8581
},
8682
Group {
8783
fn drop($self: $S::Group);
8884
fn clone($self: &$S::Group) -> $S::Group;
89-
fn new(delimiter: Delimiter, stream: $S::TokenStream) -> $S::Group;
85+
fn new(delimiter: Delimiter, stream: Option<$S::TokenStream>) -> $S::Group;
9086
fn delimiter($self: &$S::Group) -> Delimiter;
9187
fn stream($self: &$S::Group) -> $S::TokenStream;
9288
fn span($self: &$S::Group) -> $S::Span;

‎library/proc_macro/src/bridge/server.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ use super::client::HandleStore;
88
pub trait Types {
99
type FreeFunctions: 'static;
1010
type TokenStream: 'static + Clone;
11-
type TokenStreamBuilder: 'static;
12-
type TokenStreamIter: 'static + Clone;
1311
type Group: 'static + Clone;
1412
type Punct: 'static + Copy + Eq + Hash;
1513
type Ident: 'static + Copy + Eq + Hash;

0 commit comments

Comments
 (0)