Non-recursively drop TokenStream relocating only streams, not tokens

dtolnay · Jan 11, 2025 · 172907a · 172907a
1 parent acc7d36
commit 172907a
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 17 deletions.
diff --git a/src/fallback.rs b/src/fallback.rs
@@ -125,21 +125,32 @@ fn push_token_from_proc_macro(mut vec: RcVecMut<TokenTree>, token: TokenTree) {
 // Nonrecursive to prevent stack overflow.
 impl Drop for TokenStream {
     fn drop(&mut self) {
-        let mut inner = match self.inner.get_mut() {
-            Some(inner) => inner,
+        let mut stack = Vec::new();
+        let mut current = match self.inner.get_mut() {
+            Some(inner) => inner.take(),
             None => return,
         };
-        while let Some(token) = inner.pop() {
-            let group = match token {
-                TokenTree::Group(group) => group.inner,
-                _ => continue,
-            };
-            #[cfg(wrap_proc_macro)]
-            let group = match group {
-                crate::imp::Group::Fallback(group) => group,
-                crate::imp::Group::Compiler(_) => continue,
-            };
-            inner.extend(group.stream.take_inner());
+        loop {
+            while let Some(token) = current.pop() {
+                let group = match token {
+                    TokenTree::Group(group) => group.inner,
+                    _ => continue,
+                };
+                #[cfg(wrap_proc_macro)]
+                let group = match group {
+                    crate::imp::Group::Fallback(group) => group,
+                    crate::imp::Group::Compiler(_) => continue,
+                };
+                let mut group = group;
+                if let Some(inner) = group.stream.inner.get_mut() {
+                    stack.push(current);
+                    current = inner.take();
+                }
+            }
+            match stack.pop() {
+                Some(next) => current = next,
+                None => return,
+            }
         }
     }
 }

diff --git a/src/rcvec.rs b/src/rcvec.rs
@@ -80,6 +80,10 @@ impl<T> RcVecBuilder<T> {
         self.inner.extend(iter);
     }
 
+    pub(crate) fn pop(&mut self) -> Option<T> {
+        self.inner.pop()
+    }
+
     pub(crate) fn as_mut(&mut self) -> RcVecMut<T> {
         RcVecMut {
             inner: &mut self.inner,
@@ -102,13 +106,14 @@ impl<'a, T> RcVecMut<'a, T> {
         self.inner.extend(iter);
     }
 
-    pub(crate) fn pop(&mut self) -> Option<T> {
-        self.inner.pop()
-    }
-
     pub(crate) fn as_mut(&mut self) -> RcVecMut<T> {
         RcVecMut { inner: self.inner }
     }
+
+    pub(crate) fn take(self) -> RcVecBuilder<T> {
+        let vec = mem::take(self.inner);
+        RcVecBuilder { inner: vec }
+    }
 }
 
 impl<T> Clone for RcVec<T> {