diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs index 4d2049cbc41ee..1d3c4fcca0a44 100644 --- a/compiler/rustc_ast/src/tokenstream.rs +++ b/compiler/rustc_ast/src/tokenstream.rs @@ -245,12 +245,12 @@ impl AttrTokenStream { // properly implemented - we always synthesize fake tokens, // so we never reach this code. - let mut builder = TokenStreamBuilder::new(); + let mut stream = TokenStream::default(); for inner_attr in inner_attrs { - builder.push(inner_attr.tokens()); + stream.push_stream(inner_attr.tokens()); } - builder.push(delim_tokens.clone()); - *tree = TokenTree::Delimited(*span, *delim, builder.build()); + stream.push_stream(delim_tokens.clone()); + *tree = TokenTree::Delimited(*span, *delim, stream); found = true; break; } @@ -505,76 +505,49 @@ impl TokenStream { self.trees().map(|tree| TokenStream::flatten_token_tree(tree)).collect() } -} -// 99.5%+ of the time we have 1 or 2 elements in this vector. -#[derive(Clone)] -pub struct TokenStreamBuilder(SmallVec<[TokenStream; 2]>); - -impl TokenStreamBuilder { - pub fn new() -> TokenStreamBuilder { - TokenStreamBuilder(SmallVec::new()) - } - - pub fn push(&mut self, stream: TokenStream) { - self.0.push(stream); - } - - pub fn build(self) -> TokenStream { - let mut streams = self.0; - match streams.len() { - 0 => TokenStream::default(), - 1 => streams.pop().unwrap(), - _ => { - // We will extend the first stream in `streams` with the - // elements from the subsequent streams. This requires using - // `make_mut()` on the first stream, and in practice this - // doesn't cause cloning 99.9% of the time. - // - // 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. - - // Determine how much the first stream will be extended. - // Needed to avoid quadratic blow up from on-the-fly - // reallocations (#57735). - let num_appends = streams.iter().skip(1).map(|ts| ts.len()).sum(); - - // Get the first stream, which will become the result stream. - // If it's `None`, create an empty stream. - let mut iter = streams.into_iter(); - let mut res_stream_lrc = iter.next().unwrap().0; - - // Append the subsequent elements to the result stream, after - // reserving space for them. - let res_vec_mut = Lrc::make_mut(&mut res_stream_lrc); - res_vec_mut.reserve(num_appends); - for stream in iter { - let stream_iter = stream.0.iter().cloned(); - - // If (a) `res_mut_vec` is not empty and the last tree - // within it is a token tree marked with `Joint`, and (b) - // `stream` is not empty and the first tree within it is a - // token tree, and (c) the two tokens can be glued - // together... - if let Some(TokenTree::Token(last_tok, Spacing::Joint)) = res_vec_mut.last() - && let Some(TokenTree::Token(tok, spacing)) = stream.0.first() - && let Some(glued_tok) = last_tok.glue(&tok) - { - // ...then overwrite the last token tree in - // `res_vec_mut` with the glued token, and skip the - // first token tree from `stream`. - *res_vec_mut.last_mut().unwrap() = TokenTree::Token(glued_tok, *spacing); - res_vec_mut.extend(stream_iter.skip(1)); - } else { - // Append all of `stream`. - res_vec_mut.extend(stream_iter); - } - } + // If `vec` is not empty, try to glue `tt` onto its last token. The return + // value indicates if gluing took place. + fn try_glue_to_last(vec: &mut Vec, tt: &TokenTree) -> bool { + if let Some(TokenTree::Token(last_tok, Spacing::Joint)) = vec.last() + && let TokenTree::Token(tok, spacing) = tt + && let Some(glued_tok) = last_tok.glue(&tok) + { + // ...then overwrite the last token tree in `vec` with the + // glued token, and skip the first token tree from `stream`. + *vec.last_mut().unwrap() = TokenTree::Token(glued_tok, *spacing); + true + } else { + false + } + } - TokenStream(res_stream_lrc) - } + // Push `tt` onto the end of the stream, possibly gluing it to the last + // token. Uses `make_mut` to maximize efficiency. + pub fn push_tree(&mut self, tt: TokenTree) { + let vec_mut = Lrc::make_mut(&mut self.0); + + if Self::try_glue_to_last(vec_mut, &tt) { + // nothing else to do + } else { + vec_mut.push(tt); + } + } + + // Push `stream` onto the end of the stream, possibly gluing the first + // token tree to the last token. (No other token trees will be glued.) + // Uses `make_mut` to maximize efficiency. + pub fn push_stream(&mut self, stream: TokenStream) { + let vec_mut = Lrc::make_mut(&mut self.0); + + let stream_iter = stream.0.iter().cloned(); + + if let Some(first) = stream.0.first() && Self::try_glue_to_last(vec_mut, first) { + // Now skip the first token tree from `stream`. + vec_mut.extend(stream_iter.skip(1)); + } else { + // Append all of `stream`. + vec_mut.extend(stream_iter); } } } diff --git a/compiler/rustc_expand/src/proc_macro_server.rs b/compiler/rustc_expand/src/proc_macro_server.rs index 17a348ec6baae..cc2858d3f73a1 100644 --- a/compiler/rustc_expand/src/proc_macro_server.rs +++ b/compiler/rustc_expand/src/proc_macro_server.rs @@ -1,5 +1,8 @@ use crate::base::ExtCtxt; - +use pm::bridge::{ + server, DelimSpan, Diagnostic, ExpnGlobals, Group, Ident, LitKind, Literal, Punct, TokenTree, +}; +use pm::{Delimiter, Level, LineColumn}; use rustc_ast as ast; use rustc_ast::token; use rustc_ast::tokenstream::{self, Spacing::*, TokenStream}; @@ -13,11 +16,7 @@ use rustc_session::parse::ParseSess; use rustc_span::def_id::CrateNum; use rustc_span::symbol::{self, sym, Symbol}; use rustc_span::{BytePos, FileName, Pos, SourceFile, Span}; - -use pm::bridge::{ - server, DelimSpan, Diagnostic, ExpnGlobals, Group, Ident, LitKind, Literal, Punct, TokenTree, -}; -use pm::{Delimiter, Level, LineColumn}; +use smallvec::{smallvec, SmallVec}; use std::ops::Bound; trait FromInternal { @@ -253,23 +252,57 @@ impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)> for Vec for (TokenTree, &mut Rustc<'_, '_>) { - fn to_internal(self) -> TokenStream { +// We use a `SmallVec` because the output size is always one or two `TokenTree`s. +impl ToInternal> + for (TokenTree, &mut Rustc<'_, '_>) +{ + fn to_internal(self) -> SmallVec<[tokenstream::TokenTree; 2]> { use rustc_ast::token::*; let (tree, rustc) = self; - let (ch, joint, span) = match tree { - TokenTree::Punct(Punct { ch, joint, span }) => (ch, joint, span), + match tree { + TokenTree::Punct(Punct { ch, joint, span }) => { + let kind = match ch { + b'=' => Eq, + b'<' => Lt, + b'>' => Gt, + b'!' => Not, + b'~' => Tilde, + b'+' => BinOp(Plus), + b'-' => BinOp(Minus), + b'*' => BinOp(Star), + b'/' => BinOp(Slash), + b'%' => BinOp(Percent), + b'^' => BinOp(Caret), + b'&' => BinOp(And), + b'|' => BinOp(Or), + b'@' => At, + b'.' => Dot, + b',' => Comma, + b';' => Semi, + b':' => Colon, + b'#' => Pound, + b'$' => Dollar, + b'?' => Question, + b'\'' => SingleQuote, + _ => unreachable!(), + }; + smallvec![if joint { + tokenstream::TokenTree::token_joint(kind, span) + } else { + tokenstream::TokenTree::token_alone(kind, span) + }] + } TokenTree::Group(Group { delimiter, stream, span: DelimSpan { open, close, .. } }) => { - return tokenstream::TokenStream::delimited( + smallvec![tokenstream::TokenTree::Delimited( tokenstream::DelimSpan { open, close }, delimiter.to_internal(), stream.unwrap_or_default(), - ); + )] } TokenTree::Ident(self::Ident { sym, is_raw, span }) => { rustc.sess().symbol_gallery.insert(sym, span); - return tokenstream::TokenStream::token_alone(Ident(sym, is_raw), span); + smallvec![tokenstream::TokenTree::token_alone(Ident(sym, is_raw), span)] } TokenTree::Literal(self::Literal { kind: self::LitKind::Integer, @@ -282,7 +315,7 @@ impl ToInternal for (TokenTree, &mut Rus let integer = TokenKind::lit(token::Integer, symbol, suffix); let a = tokenstream::TokenTree::token_alone(minus, span); let b = tokenstream::TokenTree::token_alone(integer, span); - return [a, b].into_iter().collect(); + smallvec![a, b] } TokenTree::Literal(self::Literal { kind: self::LitKind::Float, @@ -295,46 +328,14 @@ impl ToInternal for (TokenTree, &mut Rus let float = TokenKind::lit(token::Float, symbol, suffix); let a = tokenstream::TokenTree::token_alone(minus, span); let b = tokenstream::TokenTree::token_alone(float, span); - return [a, b].into_iter().collect(); + smallvec![a, b] } TokenTree::Literal(self::Literal { kind, symbol, suffix, span }) => { - return tokenstream::TokenStream::token_alone( + smallvec![tokenstream::TokenTree::token_alone( TokenKind::lit(kind.to_internal(), symbol, suffix), span, - ); + )] } - }; - - let kind = match ch { - b'=' => Eq, - b'<' => Lt, - b'>' => Gt, - b'!' => Not, - b'~' => Tilde, - b'+' => BinOp(Plus), - b'-' => BinOp(Minus), - b'*' => BinOp(Star), - b'/' => BinOp(Slash), - b'%' => BinOp(Percent), - b'^' => BinOp(Caret), - b'&' => BinOp(And), - b'|' => BinOp(Or), - b'@' => At, - b'.' => Dot, - b',' => Comma, - b';' => Semi, - b':' => Colon, - b'#' => Pound, - b'$' => Dollar, - b'?' => Question, - b'\'' => SingleQuote, - _ => unreachable!(), - }; - - if joint { - tokenstream::TokenStream::token_joint(kind, span) - } else { - tokenstream::TokenStream::token_alone(kind, span) } } } @@ -549,7 +550,7 @@ impl server::TokenStream for Rustc<'_, '_> { &mut self, tree: TokenTree, ) -> Self::TokenStream { - (tree, &mut *self).to_internal() + Self::TokenStream::new((tree, &mut *self).to_internal().into_iter().collect::>()) } fn concat_trees( @@ -557,14 +558,14 @@ impl server::TokenStream for Rustc<'_, '_> { base: Option, trees: Vec>, ) -> Self::TokenStream { - let mut builder = tokenstream::TokenStreamBuilder::new(); - if let Some(base) = base { - builder.push(base); - } + let mut stream = + if let Some(base) = base { base } else { tokenstream::TokenStream::default() }; for tree in trees { - builder.push((tree, &mut *self).to_internal()); + for tt in (tree, &mut *self).to_internal() { + stream.push_tree(tt); + } } - builder.build() + stream } fn concat_streams( @@ -572,14 +573,12 @@ impl server::TokenStream for Rustc<'_, '_> { base: Option, streams: Vec, ) -> Self::TokenStream { - let mut builder = tokenstream::TokenStreamBuilder::new(); - if let Some(base) = base { - builder.push(base); + let mut stream = + if let Some(base) = base { base } else { tokenstream::TokenStream::default() }; + for s in streams { + stream.push_stream(s); } - for stream in streams { - builder.push(stream); - } - builder.build() + stream } fn into_trees( @@ -705,6 +704,7 @@ impl server::Span for Rustc<'_, '_> { fn source_text(&mut self, span: Self::Span) -> Option { self.sess().source_map().span_to_snippet(span).ok() } + /// Saves the provided span into the metadata of /// *the crate we are currently compiling*, which must /// be a proc-macro crate. This id can be passed to diff --git a/compiler/rustc_expand/src/tokenstream/tests.rs b/compiler/rustc_expand/src/tokenstream/tests.rs index eed69681011e3..91c4dd732e3a5 100644 --- a/compiler/rustc_expand/src/tokenstream/tests.rs +++ b/compiler/rustc_expand/src/tokenstream/tests.rs @@ -1,7 +1,7 @@ use crate::tests::string_to_stream; use rustc_ast::token; -use rustc_ast::tokenstream::{TokenStream, TokenStreamBuilder}; +use rustc_ast::tokenstream::{TokenStream, TokenTree}; use rustc_span::create_default_session_globals_then; use rustc_span::{BytePos, Span, Symbol}; @@ -19,10 +19,9 @@ fn test_concat() { let test_res = string_to_ts("foo::bar::baz"); let test_fst = string_to_ts("foo::bar"); let test_snd = string_to_ts("::baz"); - let mut builder = TokenStreamBuilder::new(); - builder.push(test_fst); - builder.push(test_snd); - let eq_res = builder.build(); + let mut eq_res = TokenStream::default(); + eq_res.push_stream(test_fst); + eq_res.push_stream(test_snd); assert_eq!(test_res.trees().count(), 5); assert_eq!(eq_res.trees().count(), 5); assert_eq!(test_res.eq_unspanned(&eq_res), true); @@ -99,11 +98,10 @@ fn test_is_empty() { #[test] fn test_dotdotdot() { create_default_session_globals_then(|| { - let mut builder = TokenStreamBuilder::new(); - builder.push(TokenStream::token_joint(token::Dot, sp(0, 1))); - builder.push(TokenStream::token_joint(token::Dot, sp(1, 2))); - builder.push(TokenStream::token_alone(token::Dot, sp(2, 3))); - let stream = builder.build(); + let mut stream = TokenStream::default(); + stream.push_tree(TokenTree::token_joint(token::Dot, sp(0, 1))); + stream.push_tree(TokenTree::token_joint(token::Dot, sp(1, 2))); + stream.push_tree(TokenTree::token_alone(token::Dot, sp(2, 3))); assert!(stream.eq_unspanned(&string_to_ts("..."))); assert_eq!(stream.trees().count(), 1); }) diff --git a/library/proc_macro/src/bridge/client.rs b/library/proc_macro/src/bridge/client.rs index 4461b21802adb..506b2a773cc49 100644 --- a/library/proc_macro/src/bridge/client.rs +++ b/library/proc_macro/src/bridge/client.rs @@ -223,10 +223,10 @@ pub(crate) use super::symbol::Symbol; macro_rules! define_client_side { ($($name:ident { - $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)* + $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)* }),* $(,)?) => { $(impl $name { - $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)* { + $(pub(crate) fn $method($($arg: $arg_ty),*) $(-> $ret_ty)? { Bridge::with(|bridge| { let mut buf = bridge.cached_buffer.take();