diff --git a/src/iter.rs b/src/iter.rs index c25be0b..722013c 100644 --- a/src/iter.rs +++ b/src/iter.rs @@ -1,21 +1,23 @@ -use proc_macro::{token_stream, TokenStream, TokenTree}; -use std::iter::Peekable; +use proc_macro::{token_stream, Delimiter, TokenStream, TokenTree}; pub type Iter<'a> = &'a mut IterImpl; pub struct IterImpl { - tokens: Peekable, + stack: Vec, + peeked: Option, } pub fn new(tokens: TokenStream) -> IterImpl { IterImpl { - tokens: tokens.into_iter().peekable(), + stack: vec![tokens.into_iter()], + peeked: None, } } impl IterImpl { pub fn peek(&mut self) -> Option<&TokenTree> { - self.tokens.peek() + self.peeked = self.next(); + self.peeked.as_ref() } } @@ -23,6 +25,18 @@ impl Iterator for IterImpl { type Item = TokenTree; fn next(&mut self) -> Option { - self.tokens.next() + if let Some(tt) = self.peeked.take() { + return Some(tt); + } + loop { + let top = self.stack.last_mut()?; + match top.next() { + None => drop(self.stack.pop()), + Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::None => { + self.stack.push(group.stream().into_iter()); + } + Some(tt) => return Some(tt), + } + } } }