1 files changed, 50 insertions, 1 deletions

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index 84cb417dd..28fcf80b3 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 }