3 files changed, 55 insertions, 8 deletions

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index 84cb417dd89..28fcf80b31b 100644
--- a/compiler/rustc_data_structures/src/functor.rs
+++ b/compiler/rustc_data_structures/src/functor.rs
@@ -1,5 +1,5 @@
 use rustc_index::vec::{Idx, IndexVec};
-use std::mem;
+use std::{mem, rc::Rc, sync::Arc};
 
 pub trait IdFunctor: Sized {
     type Inner;
@@ -65,3 +65,52 @@ impl<I: Idx, T> IdFunctor for IndexVec<I, T> {
         self.raw.try_map_id(f).map(IndexVec::from_raw)
     }
 }
+
+macro_rules! rc {
+    ($($rc:ident),+) => {$(
+        impl<T: Clone> IdFunctor for $rc<T> {
+            type Inner = T;
+
+            #[inline]
+            fn try_map_id<F, E>(mut self, mut f: F) -> Result<Self, E>
+            where
+                F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
+            {
+                // We merely want to replace the contained `T`, if at all possible,
+                // so that we don't needlessly allocate a new `$rc` or indeed clone
+                // the contained type.
+                unsafe {
+                    // First step is to ensure that we have a unique reference to
+                    // the contained type, which `$rc::make_mut` will accomplish (by
+                    // allocating a new `$rc` and cloning the `T` only if required).
+                    // This is done *before* casting to `$rc<ManuallyDrop<T>>` so that
+                    // panicking during `make_mut` does not leak the `T`.
+                    $rc::make_mut(&mut self);
+
+                    // Casting to `$rc<ManuallyDrop<T>>` is safe because `ManuallyDrop`
+                    // is `repr(transparent)`.
+                    let ptr = $rc::into_raw(self).cast::<mem::ManuallyDrop<T>>();
+                    let mut unique = $rc::from_raw(ptr);
+
+                    // Call to `$rc::make_mut` above guarantees that `unique` is the
+                    // sole reference to the contained value, so we can avoid doing
+                    // a checked `get_mut` here.
+                    let slot = $rc::get_mut_unchecked(&mut unique);
+
+                    // Semantically move the contained type out from `unique`, fold
+                    // it, then move the folded value back into `unique`. Should
+                    // folding fail, `ManuallyDrop` ensures that the "moved-out"
+                    // value is not re-dropped.
+                    let owned = mem::ManuallyDrop::take(slot);
+                    let folded = f(owned)?;
+                    *slot = mem::ManuallyDrop::new(folded);
+
+                    // Cast back to `$rc<T>`.
+                    Ok($rc::from_raw($rc::into_raw(unique).cast()))
+                }
+            }
+        }
+    )+};
+}
+
+rc! { Rc, Arc }
diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 7fab8954cb1..a94e52fdfe6 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -26,6 +26,7 @@
 #![feature(test)]
 #![feature(thread_id_value)]
 #![feature(vec_into_raw_parts)]
+#![feature(get_mut_unchecked)]
 #![allow(rustc::default_hash_types)]
 #![allow(rustc::potential_query_instability)]
 #![deny(rustc::untranslatable_diagnostic)]
diff --git a/compiler/rustc_data_structures/src/profiling.rs b/compiler/rustc_data_structures/src/profiling.rs
index 3aca03f6e5c..44331683694 100644
--- a/compiler/rustc_data_structures/src/profiling.rs
+++ b/compiler/rustc_data_structures/src/profiling.rs
@@ -207,8 +207,7 @@ impl SelfProfilerRef {
     /// a measureme event, "verbose" generic activities also print a timing entry to
     /// stderr if the compiler is invoked with -Ztime-passes.
     pub fn verbose_generic_activity(&self, event_label: &'static str) -> VerboseTimingGuard<'_> {
-        let message =
-            if self.print_verbose_generic_activities { Some(event_label.to_owned()) } else { None };
+        let message = self.print_verbose_generic_activities.then(|| event_label.to_owned());
 
         VerboseTimingGuard::start(message, self.generic_activity(event_label))
     }
@@ -222,11 +221,9 @@ impl SelfProfilerRef {
     where
         A: Borrow<str> + Into<String>,
     {
-        let message = if self.print_verbose_generic_activities {
-            Some(format!("{}({})", event_label, event_arg.borrow()))
-        } else {
-            None
-        };
+        let message = self
+            .print_verbose_generic_activities
+            .then(|| format!("{}({})", event_label, event_arg.borrow()));
 
         VerboseTimingGuard::start(message, self.generic_activity_with_arg(event_label, event_arg))
     }