From 14e57d7c134cc80878c10f2e2bb61e62cc6a2fdb Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Sun, 2 Jul 2023 22:43:39 +0200
Subject: [PATCH 1/2] perform TokenStream replacement in-place when possible in
 expand_macro

---
 compiler/rustc_ast/src/tokenstream.rs        | 21 +++++++++++++++++---
 compiler/rustc_expand/src/mbe/macro_rules.rs |  3 +--
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
index db296aa44db2b..45c0586424a3b 100644
--- a/compiler/rustc_ast/src/tokenstream.rs
+++ b/compiler/rustc_ast/src/tokenstream.rs
@@ -25,7 +25,7 @@ use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
 use rustc_span::{Span, DUMMY_SP};
 use smallvec::{smallvec, SmallVec};
 
-use std::{fmt, iter};
+use std::{fmt, iter, mem};
 
 /// When the main Rust parser encounters a syntax-extension invocation, it
 /// parses the arguments to the invocation as a token tree. This is a very
@@ -410,8 +410,23 @@ impl TokenStream {
         t1.next().is_none() && t2.next().is_none()
     }
 
-    pub fn map_enumerated<F: FnMut(usize, &TokenTree) -> TokenTree>(self, mut f: F) -> TokenStream {
-        TokenStream(Lrc::new(self.0.iter().enumerate().map(|(i, tree)| f(i, tree)).collect()))
+    /// Applies the supplied function to each `TokenTree` and its index in `self`, returning a new `TokenStream`
+    ///
+    /// It is equivalent to `TokenStream::new(self.trees().cloned().enumerate().map(|(i, tt)| f(i, tt)).collect())`.
+    pub fn map_enumerated_owned(
+        mut self,
+        mut f: impl FnMut(usize, TokenTree) -> TokenTree,
+    ) -> TokenStream {
+        if let Some(inner) = Lrc::get_mut(&mut self.0) {
+            // optimization: perform the map in-place if self's reference count is 1
+            let owned = mem::take(inner);
+            *inner = owned.into_iter().enumerate().map(|(i, tree)| f(i, tree)).collect();
+            self
+        } else {
+            TokenStream(Lrc::new(
+                self.0.iter().enumerate().map(|(i, tree)| f(i, tree.clone())).collect(),
+            ))
+        }
     }
 
     /// Create a token stream containing a single token with alone spacing.
diff --git a/compiler/rustc_expand/src/mbe/macro_rules.rs b/compiler/rustc_expand/src/mbe/macro_rules.rs
index 4a8a14994ff11..42cc0a6b143ce 100644
--- a/compiler/rustc_expand/src/mbe/macro_rules.rs
+++ b/compiler/rustc_expand/src/mbe/macro_rules.rs
@@ -223,8 +223,7 @@ fn expand_macro<'cx>(
             // Replace all the tokens for the corresponding positions in the macro, to maintain
             // proper positions in error reporting, while maintaining the macro_backtrace.
             if tts.len() == rhs.tts.len() {
-                tts = tts.map_enumerated(|i, tt| {
-                    let mut tt = tt.clone();
+                tts = tts.map_enumerated_owned(|i, mut tt| {
                     let rhs_tt = &rhs.tts[i];
                     let ctxt = tt.span().ctxt();
                     match (&mut tt, rhs_tt) {

From 7916a2c1c9eda7e91ce1ed4c714f7d12daf13221 Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Mon, 3 Jul 2023 13:40:25 +0200
Subject: [PATCH 2/2] Use Lrc::make_mut instead of Lrc::get_mut

---
 compiler/rustc_ast/src/tokenstream.rs | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
index 45c0586424a3b..ca4a739abd7a1 100644
--- a/compiler/rustc_ast/src/tokenstream.rs
+++ b/compiler/rustc_ast/src/tokenstream.rs
@@ -417,16 +417,10 @@ impl TokenStream {
         mut self,
         mut f: impl FnMut(usize, TokenTree) -> TokenTree,
     ) -> TokenStream {
-        if let Some(inner) = Lrc::get_mut(&mut self.0) {
-            // optimization: perform the map in-place if self's reference count is 1
-            let owned = mem::take(inner);
-            *inner = owned.into_iter().enumerate().map(|(i, tree)| f(i, tree)).collect();
-            self
-        } else {
-            TokenStream(Lrc::new(
-                self.0.iter().enumerate().map(|(i, tree)| f(i, tree.clone())).collect(),
-            ))
-        }
+        let owned = Lrc::make_mut(&mut self.0); // clone if necessary
+        // rely on vec's in-place optimizations to avoid another allocation
+        *owned = mem::take(owned).into_iter().enumerate().map(|(i, tree)| f(i, tree)).collect();
+        self
     }
 
     /// Create a token stream containing a single token with alone spacing.