Auto merge of #91353 - eggyal:reuse-rcs-during-folding, r=lcnr

Avoid cloning refcounted types during folding

Addresses FIXME comment created in #78313

r? `@lcnr`

2 files changed, 68 insertions, 6 deletions

diff --git a/compiler/rustc_middle/src/lib.rs b/compiler/rustc_middle/src/lib.rs
index 66d1ae1420a..9d8588c9a8e 100644
--- a/compiler/rustc_middle/src/lib.rs
+++ b/compiler/rustc_middle/src/lib.rs
@@ -33,6 +33,7 @@
 #![feature(derive_default_enum)]
 #![feature(discriminant_kind)]
 #![feature(exhaustive_patterns)]
+#![feature(get_mut_unchecked)]
 #![feature(if_let_guard)]
 #![feature(map_first_last)]
 #![feature(never_type)]
diff --git a/compiler/rustc_middle/src/ty/structural_impls.rs b/compiler/rustc_middle/src/ty/structural_impls.rs
index 30ed008a5de..28dc9767b78 100644
--- a/compiler/rustc_middle/src/ty/structural_impls.rs
+++ b/compiler/rustc_middle/src/ty/structural_impls.rs
@@ -13,6 +13,7 @@ use rustc_hir::def_id::CRATE_DEF_INDEX;
 use rustc_index::vec::{Idx, IndexVec};
 
 use std::fmt;
+use std::mem::ManuallyDrop;
 use std::ops::ControlFlow;
 use std::rc::Rc;
 use std::sync::Arc;
@@ -732,11 +733,41 @@ EnumTypeFoldableImpl! {
 
 impl<'tcx, T: TypeFoldable<'tcx>> TypeFoldable<'tcx> for Rc<T> {
     fn try_super_fold_with<F: FallibleTypeFolder<'tcx>>(
-        self,
+        mut self,
         folder: &mut F,
     ) -> Result<Self, F::Error> {
-        // FIXME: Reuse the `Rc` here.
-        (*self).clone().try_fold_with(folder).map(Rc::new)
+        // 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::<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 = ManuallyDrop::take(slot);
+            let folded = owned.try_fold_with(folder)?;
+            *slot = ManuallyDrop::new(folded);
+
+            // Cast back to `Rc<T>`.
+            Ok(Rc::from_raw(Rc::into_raw(unique).cast()))
+        }
     }
 
     fn super_visit_with<V: TypeVisitor<'tcx>>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
@@ -746,11 +777,41 @@ impl<'tcx, T: TypeFoldable<'tcx>> TypeFoldable<'tcx> for Rc<T> {
 
 impl<'tcx, T: TypeFoldable<'tcx>> TypeFoldable<'tcx> for Arc<T> {
     fn try_super_fold_with<F: FallibleTypeFolder<'tcx>>(
-        self,
+        mut self,
         folder: &mut F,
     ) -> Result<Self, F::Error> {
-        // FIXME: Reuse the `Arc` here.
-        (*self).clone().try_fold_with(folder).map(Arc::new)
+        // We merely want to replace the contained `T`, if at all possible,
+        // so that we don't needlessly allocate a new `Arc` or indeed clone
+        // the contained type.
+        unsafe {
+            // First step is to ensure that we have a unique reference to
+            // the contained type, which `Arc::make_mut` will accomplish (by
+            // allocating a new `Arc` and cloning the `T` only if required).
+            // This is done *before* casting to `Arc<ManuallyDrop<T>>` so that
+            // panicking during `make_mut` does not leak the `T`.
+            Arc::make_mut(&mut self);
+
+            // Casting to `Arc<ManuallyDrop<T>>` is safe because `ManuallyDrop`
+            // is `repr(transparent)`.
+            let ptr = Arc::into_raw(self).cast::<ManuallyDrop<T>>();
+            let mut unique = Arc::from_raw(ptr);
+
+            // Call to `Arc::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 = Arc::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 = ManuallyDrop::take(slot);
+            let folded = owned.try_fold_with(folder)?;
+            *slot = ManuallyDrop::new(folded);
+
+            // Cast back to `Arc<T>`.
+            Ok(Arc::from_raw(Arc::into_raw(unique).cast()))
+        }
     }
 
     fn super_visit_with<V: TypeVisitor<'tcx>>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {