Redo TokenStreamIter

compiler/rustc_ast/src/tokenstream.rs

@@ -22,6 +22,9 @@ use crate::ast_traits::{HasAttrs, HasTokens};
 use crate::token::{self, Delimiter, Token, TokenKind};
 use crate::{AttrVec, Attribute};
 
+#[cfg(test)]
+mod tests;
+
 /// Part of a `TokenStream`.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Encodable, Decodable, StableHash)]
 pub enum TokenTree {
@@ -833,32 +836,30 @@ impl StableHash for TokenStream {
 }
 
 #[derive(Clone)]
-pub struct TokenStreamIter<'t> {
-    stream: &'t TokenStream,
-    index: usize,
-}
+pub struct TokenStreamIter<'t>(std::slice::Iter<'t, TokenTree>);
 
 impl<'t> TokenStreamIter<'t> {
     fn new(stream: &'t TokenStream) -> Self {
-        TokenStreamIter { stream, index: 0 }
+        TokenStreamIter(stream.0.as_slice().iter())
     }
 
     // Peeking could be done via `Peekable`, but most iterators need peeking,
     // and this is simple and avoids the need to use `peekable` and `Peekable`
     // at all the use sites.
     pub fn peek(&self) -> Option<&'t TokenTree> {
-        self.stream.0.get(self.index)
+        self.0.as_slice().first()
     }
 }
 
 impl<'t> Iterator for TokenStreamIter<'t> {
     type Item = &'t TokenTree;
 
     fn next(&mut self) -> Option<&'t TokenTree> {
-        self.stream.0.get(self.index).map(|tree| {
-            self.index += 1;
-            tree
-        })
+        self.0.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
     }
 }
 

compiler/rustc_ast/src/tokenstream/tests.rs

@@ -0,0 +1,13 @@
+use rustc_span::DUMMY_SP;
+
+use crate::token::TokenKind;
+use crate::tokenstream::TokenStream;
+
+#[test]
+fn test_token_stream_iter() {
+    let ts = TokenStream::token_alone(TokenKind::Eq, DUMMY_SP);
+    assert_eq!(ts.len(), 1);
+
+    let iter = ts.iter();
+    assert_eq!(iter.size_hint(), (1, Some(1)));
+}