Remove allocation from RcVec -> RcVecBuilder conversion

dtolnay · Jul 25, 2022 · a2b46f9 · a2b46f9
1 parent d8c056b
commit a2b46f9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/src/fallback.rs b/src/fallback.rs
@@ -7,10 +7,12 @@ use std::cell::RefCell;
 use std::cmp;
 use std::fmt::{self, Debug, Display, Write};
 use std::iter::FromIterator;
+use std::mem::ManuallyDrop;
 use std::ops::RangeBounds;
 #[cfg(procmacro2_semver_exempt)]
 use std::path::Path;
 use std::path::PathBuf;
+use std::ptr;
 use std::str::FromStr;
 
 /// Force use of proc-macro2's fallback implementation of the API for now, even
@@ -60,8 +62,9 @@ impl TokenStream {
         self.inner.len() == 0
     }
 
-    fn take_inner(&mut self) -> RcVecBuilder<TokenTree> {
-        self.inner.make_mut().take()
+    fn take_inner(self) -> RcVecBuilder<TokenTree> {
+        let nodrop = ManuallyDrop::new(self);
+        unsafe { ptr::read(&nodrop.inner) }.make_owned()
     }
 }
 
@@ -122,7 +125,6 @@ impl Drop for TokenStream {
                 crate::imp::Group::Fallback(group) => group,
                 crate::imp::Group::Compiler(_) => continue,
             };
-            let mut group = group;
             inner.extend(group.stream.take_inner());
         }
     }
@@ -264,7 +266,7 @@ impl FromIterator<TokenStream> for TokenStream {
     fn from_iter<I: IntoIterator<Item = TokenStream>>(streams: I) -> Self {
         let mut v = RcVecBuilder::new();
 
-        for mut stream in streams {
+        for stream in streams {
             v.extend(stream.take_inner());
         }
 
@@ -293,7 +295,7 @@ impl IntoIterator for TokenStream {
     type Item = TokenTree;
     type IntoIter = TokenTreeIter;
 
-    fn into_iter(mut self) -> TokenTreeIter {
+    fn into_iter(self) -> TokenTreeIter {
         self.take_inner().into_iter()
     }
 }

diff --git a/src/rcvec.rs b/src/rcvec.rs
@@ -46,6 +46,18 @@ impl<T> RcVec<T> {
         let inner = Rc::get_mut(&mut self.inner)?;
         Some(RcVecMut { inner })
     }
+
+    pub fn make_owned(mut self) -> RcVecBuilder<T>
+    where
+        T: Clone,
+    {
+        let vec = if let Some(owned) = Rc::get_mut(&mut self.inner) {
+            mem::replace(owned, Vec::new())
+        } else {
+            Vec::clone(&self.inner)
+        };
+        RcVecBuilder { inner: vec }
+    }
 }
 
 impl<T> RcVecBuilder<T> {
@@ -96,12 +108,6 @@ impl<'a, T> RcVecMut<'a, T> {
     pub fn as_mut(&mut self) -> RcVecMut<T> {
         RcVecMut { inner: self.inner }
     }
-
-    pub fn take(self) -> RcVecBuilder<T> {
-        RcVecBuilder {
-            inner: mem::replace(self.inner, Vec::new()),
-        }
-    }
 }
 
 impl<T> Clone for RcVec<T> {