Skip to content

Commit b0ca197

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 8db4f28 commit b0ca197

File tree

5 files changed

+130
-97
lines changed

5 files changed

+130
-97
lines changed

‎compiler/rustc_expand/src/proc_macro_server.rs

Lines changed: 59 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -284,12 +284,6 @@ impl ToInternal<rustc_errors::Level> for Level {
284284

285285
pub struct FreeFunctions;
286286

287-
#[derive(Clone)]
288-
pub struct TokenStreamIter {
289-
cursor: tokenstream::Cursor,
290-
stack: Vec<TokenTree<Group, Punct, Ident, Literal>>,
291-
}
292-
293287
#[derive(Clone)]
294288
pub struct Group {
295289
delimiter: Delimiter,
@@ -398,8 +392,6 @@ impl<'a> Rustc<'a> {
398392
impl server::Types for Rustc<'_> {
399393
type FreeFunctions = FreeFunctions;
400394
type TokenStream = TokenStream;
401-
type TokenStreamBuilder = tokenstream::TokenStreamBuilder;
402-
type TokenStreamIter = TokenStreamIter;
403395
type Group = Group;
404396
type Punct = Punct;
405397
type Ident = Ident;
@@ -417,9 +409,6 @@ impl server::FreeFunctions for Rustc<'_> {
417409
}
418410

419411
impl server::TokenStream for Rustc<'_> {
420-
fn new(&mut self) -> Self::TokenStream {
421-
TokenStream::default()
422-
}
423412
fn is_empty(&mut self, stream: &Self::TokenStream) -> bool {
424413
stream.is_empty()
425414
}
@@ -440,53 +429,74 @@ impl server::TokenStream for Rustc<'_> {
440429
) -> Self::TokenStream {
441430
tree.to_internal()
442431
}
443-
fn into_iter(&mut self, stream: Self::TokenStream) -> Self::TokenStreamIter {
444-
TokenStreamIter { cursor: stream.trees(), stack: vec![] }
445-
}
446-
}
447-
448-
impl server::TokenStreamBuilder for Rustc<'_> {
449-
fn new(&mut self) -> Self::TokenStreamBuilder {
450-
tokenstream::TokenStreamBuilder::new()
451-
}
452-
fn push(&mut self, builder: &mut Self::TokenStreamBuilder, stream: Self::TokenStream) {
453-
builder.push(stream);
432+
fn concat_trees(
433+
&mut self,
434+
base: Option<Self::TokenStream>,
435+
trees: Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>>,
436+
) -> Self::TokenStream {
437+
let mut builder = tokenstream::TokenStreamBuilder::new();
438+
if let Some(base) = base {
439+
builder.push(base);
440+
}
441+
for tree in trees {
442+
builder.push(tree.to_internal());
443+
}
444+
builder.build()
454445
}
455-
fn build(&mut self, builder: Self::TokenStreamBuilder) -> Self::TokenStream {
446+
fn concat_streams(
447+
&mut self,
448+
base: Option<Self::TokenStream>,
449+
streams: Vec<Self::TokenStream>,
450+
) -> Self::TokenStream {
451+
let mut builder = tokenstream::TokenStreamBuilder::new();
452+
if let Some(base) = base {
453+
builder.push(base);
454+
}
455+
for stream in streams {
456+
builder.push(stream);
457+
}
456458
builder.build()
457459
}
458-
}
459-
460-
impl server::TokenStreamIter for Rustc<'_> {
461-
fn next(
460+
fn into_iter(
462461
&mut self,
463-
iter: &mut Self::TokenStreamIter,
464-
) -> Option<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
462+
stream: Self::TokenStream,
463+
) -> Vec<TokenTree<Self::Group, Self::Punct, Self::Ident, Self::Literal>> {
464+
// XXX: This is a raw port of the previous approach, and can probably be
465+
// optimized.
466+
let mut cursor = stream.trees();
467+
let mut stack = Vec::new();
468+
let mut tts = Vec::new();
465469
loop {
466-
let tree = iter.stack.pop().or_else(|| {
467-
let next = iter.cursor.next_with_spacing()?;
468-
Some(TokenTree::from_internal((next, &mut iter.stack, self)))
469-
})?;
470-
// A hack used to pass AST fragments to attribute and derive macros
471-
// as a single nonterminal token instead of a token stream.
472-
// Such token needs to be "unwrapped" and not represented as a delimited group.
473-
// FIXME: It needs to be removed, but there are some compatibility issues (see #73345).
474-
if let TokenTree::Group(ref group) = tree {
475-
if group.flatten {
476-
iter.cursor.append(group.stream.clone());
477-
continue;
470+
let next = stack.pop().or_else(|| {
471+
let next = cursor.next_with_spacing()?;
472+
Some(TokenTree::from_internal((next, &mut stack, self)))
473+
});
474+
match next {
475+
Some(TokenTree::Group(group)) => {
476+
// A hack used to pass AST fragments to attribute and derive
477+
// macros as a single nonterminal token instead of a token
478+
// stream. Such token needs to be "unwrapped" and not
479+
// represented as a delimited group.
480+
// FIXME: It needs to be removed, but there are some
481+
// compatibility issues (see #73345).
482+
if group.flatten {
483+
cursor.append(group.stream);
484+
continue;
485+
}
486+
tts.push(TokenTree::Group(group));
478487
}
488+
Some(tt) => tts.push(tt),
489+
None => return tts,
479490
}
480-
return Some(tree);
481491
}
482492
}
483493
}
484494

485495
impl server::Group for Rustc<'_> {
486-
fn new(&mut self, delimiter: Delimiter, stream: Self::TokenStream) -> Self::Group {
496+
fn new(&mut self, delimiter: Delimiter, stream: Option<Self::TokenStream>) -> Self::Group {
487497
Group {
488498
delimiter,
489-
stream,
499+
stream: stream.unwrap_or_default(),
490500
span: DelimSpan::from_single(server::Context::call_site(self)),
491501
flatten: false,
492502
}
@@ -519,7 +529,11 @@ impl server::Punct for Rustc<'_> {
519529
punct.ch
520530
}
521531
fn spacing(&mut self, punct: Self::Punct) -> Spacing {
522-
if punct.joint { Spacing::Joint } else { Spacing::Alone }
532+
if punct.joint {
533+
Spacing::Joint
534+
} else {
535+
Spacing::Alone
536+
}
523537
}
524538
fn span(&mut self, punct: Self::Punct) -> Self::Span {
525539
punct.span

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

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -195,8 +195,6 @@ define_handles! {
195195
'owned:
196196
FreeFunctions,
197197
TokenStream,
198-
TokenStreamBuilder,
199-
TokenStreamIter,
200198
Group,
201199
Literal,
202200
SourceFile,
@@ -221,12 +219,6 @@ impl Clone for TokenStream {
221219
}
222220
}
223221

224-
impl Clone for TokenStreamIter {
225-
fn clone(&self) -> Self {
226-
self.clone()
227-
}
228-
}
229-
230222
impl Clone for Group {
231223
fn clone(&self) -> Self {
232224
self.clone()
@@ -507,7 +499,11 @@ impl Client<fn(crate::TokenStream) -> crate::TokenStream> {
507499
bridge: BridgeConfig<'_>,
508500
f: impl FnOnce(crate::TokenStream) -> crate::TokenStream,
509501
) -> Buffer<u8> {
510-
run_client(bridge, |input| f(crate::TokenStream(input)).0)
502+
run_client(bridge, |input| {
503+
f(crate::TokenStream(Some(input)))
504+
.0
505+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
506+
})
511507
}
512508
Client { get_handle_counters: HandleCounters::get, run, f }
513509
}
@@ -523,7 +519,9 @@ impl Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream> {
523519
f: impl FnOnce(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
524520
) -> Buffer<u8> {
525521
run_client(bridge, |(input, input2)| {
526-
f(crate::TokenStream(input), crate::TokenStream(input2)).0
522+
f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2)))
523+
.0
524+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
527525
})
528526
}
529527
Client { get_handle_counters: HandleCounters::get, run, f }

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

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -65,32 +65,28 @@ macro_rules! with_api {
6565
TokenStream {
6666
nowait fn drop($self: $S::TokenStream);
6767
nowait fn clone($self: &$S::TokenStream) -> $S::TokenStream;
68-
nowait fn new() -> $S::TokenStream;
6968
wait fn is_empty($self: &$S::TokenStream) -> bool;
7069
wait fn from_str(src: &str) -> $S::TokenStream;
7170
wait fn to_string($self: &$S::TokenStream) -> String;
7271
nowait fn from_token_tree(
7372
tree: TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>,
7473
) -> $S::TokenStream;
75-
nowait fn into_iter($self: $S::TokenStream) -> $S::TokenStreamIter;
76-
},
77-
TokenStreamBuilder {
78-
nowait fn drop($self: $S::TokenStreamBuilder);
79-
nowait fn new() -> $S::TokenStreamBuilder;
80-
nowait fn push($self: &mut $S::TokenStreamBuilder, stream: $S::TokenStream);
81-
nowait fn build($self: $S::TokenStreamBuilder) -> $S::TokenStream;
82-
},
83-
TokenStreamIter {
84-
nowait fn drop($self: $S::TokenStreamIter);
85-
nowait fn clone($self: &$S::TokenStreamIter) -> $S::TokenStreamIter;
86-
wait fn next(
87-
$self: &mut $S::TokenStreamIter,
88-
) -> Option<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
74+
nowait fn concat_trees(
75+
base: Option<$S::TokenStream>,
76+
trees: Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>,
77+
) -> $S::TokenStream;
78+
nowait fn concat_streams(
79+
base: Option<$S::TokenStream>,
80+
trees: Vec<$S::TokenStream>,
81+
) -> $S::TokenStream;
82+
wait fn into_iter(
83+
$self: $S::TokenStream
84+
) -> Vec<TokenTree<$S::Group, $S::Punct, $S::Ident, $S::Literal>>;
8985
},
9086
Group {
9187
nowait fn drop($self: $S::Group);
9288
nowait fn clone($self: &$S::Group) -> $S::Group;
93-
nowait fn new(delimiter: Delimiter, stream: $S::TokenStream) -> $S::Group;
89+
nowait fn new(delimiter: Delimiter, stream: Option<$S::TokenStream>) -> $S::Group;
9490
wait fn delimiter($self: &$S::Group) -> Delimiter;
9591
nowait fn stream($self: &$S::Group) -> $S::TokenStream;
9692
wait fn span($self: &$S::Group) -> $S::Span;

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,6 @@ macro_rules! associated_item {
1414
(type FreeFunctions: 'static;);
1515
(type TokenStream) =>
1616
(type TokenStream: 'static + Clone;);
17-
(type TokenStreamBuilder) =>
18-
(type TokenStreamBuilder: 'static;);
19-
(type TokenStreamIter) =>
20-
(type TokenStreamIter: 'static + Clone;);
2117
(type Group) =>
2218
(type Group: 'static + Clone;);
2319
(type Punct) =>

0 commit comments

Comments
 (0)