From 51a988bc02d8f872e8d77bdb70848e3a27de00ba Mon Sep 17 00:00:00 2001
From: David Tolnay <dtolnay@gmail.com>
Date: Sun, 24 May 2020 15:15:10 -0700
Subject: [PATCH] Traverse into None-delimited groups

---
 src/iter.rs | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

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<token_stream::IntoIter>,
+    stack: Vec<token_stream::IntoIter>,
+    peeked: Option<TokenTree>,
 }
 
 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::Item> {
-        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),
+            }
+        }
     }
 }