Skip to content

Commit 29d8b9c

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 8315ce9 commit 29d8b9c

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,
@@ -385,8 +379,6 @@ impl<'a, 'b> Rustc<'a, 'b> {
385379
impl server::Types for Rustc<'_, '_> {
386380
type FreeFunctions = FreeFunctions;
387381
type TokenStream = TokenStream;
388-
type TokenStreamBuilder = tokenstream::TokenStreamBuilder;
389-
type TokenStreamIter = TokenStreamIter;
390382
type Group = Group;
391383
type Punct = Punct;
392384
type Ident = Ident;
@@ -411,9 +403,6 @@ impl server::FreeFunctions for Rustc<'_, '_> {
411403
}
412404

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

529539
impl server::Group for Rustc<'_, '_> {
530-
fn new(&mut self, delimiter: Delimiter, stream: Self::TokenStream) -> Self::Group {
540+
fn new(&mut self, delimiter: Delimiter, stream: Option<Self::TokenStream>) -> Self::Group {
531541
Group {
532542
delimiter,
533-
stream,
543+
stream: stream.unwrap_or_default(),
534544
span: DelimSpan::from_single(server::Context::call_site(self)),
535545
flatten: false,
536546
}

‎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()
@@ -465,7 +457,11 @@ impl Client<crate::TokenStream, crate::TokenStream> {
465457
Client {
466458
get_handle_counters: HandleCounters::get,
467459
run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
468-
run_client(bridge, |input| f(crate::TokenStream(input)).0)
460+
run_client(bridge, |input| {
461+
f(crate::TokenStream(Some(input)))
462+
.0
463+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
464+
})
469465
}),
470466
_marker: PhantomData,
471467
}
@@ -480,7 +476,9 @@ impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
480476
get_handle_counters: HandleCounters::get,
481477
run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
482478
run_client(bridge, |(input, input2)| {
483-
f(crate::TokenStream(input), crate::TokenStream(input2)).0
479+
f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2)))
480+
.0
481+
.unwrap_or_else(|| TokenStream::concat_streams(None, vec![]))
484482
})
485483
}),
486484
_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)