@@ -25,7 +25,7 @@ use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
2525use rustc_span:: { Span , DUMMY_SP } ;
2626use smallvec:: { smallvec, SmallVec } ;
2727
28- use std:: { fmt, iter} ;
28+ use std:: { fmt, iter, mem } ;
2929
3030/// When the main Rust parser encounters a syntax-extension invocation, it
3131/// parses the arguments to the invocation as a token tree. This is a very
@@ -410,8 +410,17 @@ impl TokenStream {
410410 t1. next ( ) . is_none ( ) && t2. next ( ) . is_none ( )
411411 }
412412
413- pub fn map_enumerated < F : FnMut ( usize , & TokenTree ) -> TokenTree > ( self , mut f : F ) -> TokenStream {
414- TokenStream ( Lrc :: new ( self . 0 . iter ( ) . enumerate ( ) . map ( |( i, tree) | f ( i, tree) ) . collect ( ) ) )
413+ /// Applies the supplied function to each `TokenTree` and its index in `self`, returning a new `TokenStream`
414+ ///
415+ /// It is equivalent to `TokenStream::new(self.trees().cloned().enumerate().map(|(i, tt)| f(i, tt)).collect())`.
416+ pub fn map_enumerated_owned (
417+ mut self ,
418+ mut f : impl FnMut ( usize , TokenTree ) -> TokenTree ,
419+ ) -> TokenStream {
420+ let owned = Lrc :: make_mut ( & mut self . 0 ) ; // clone if necessary
421+ // rely on vec's in-place optimizations to avoid another allocation
422+ * owned = mem:: take ( owned) . into_iter ( ) . enumerate ( ) . map ( |( i, tree) | f ( i, tree) ) . collect ( ) ;
423+ self
415424 }
416425
417426 /// Create a token stream containing a single token with alone spacing.
0 commit comments