Refactor refcounted structural_impls via functors

4 files changed, 55 insertions, 69 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_type_ir/src/lib.rs b/compiler/rustc_type_ir/src/lib.rs
index 29d261fda8d..9db1b0fd13d 100644
--- a/compiler/rustc_type_ir/src/lib.rs
+++ b/compiler/rustc_type_ir/src/lib.rs
@@ -1,6 +1,5 @@
 #![feature(associated_type_defaults)]
 #![feature(fmt_helpers_for_derive)]
-#![feature(get_mut_unchecked)]
 #![feature(min_specialization)]
 #![feature(never_type)]
 #![feature(rustc_attrs)]
diff --git a/compiler/rustc_type_ir/src/structural_impls.rs b/compiler/rustc_type_ir/src/structural_impls.rs
index b6abe9a7357..3ebe241042f 100644
--- a/compiler/rustc_type_ir/src/structural_impls.rs
+++ b/compiler/rustc_type_ir/src/structural_impls.rs
@@ -8,7 +8,6 @@ use crate::Interner;
 use rustc_data_structures::functor::IdFunctor;
 use rustc_index::vec::{Idx, IndexVec};
 
-use std::mem::ManuallyDrop;
 use std::ops::ControlFlow;
 use std::rc::Rc;
 use std::sync::Arc;
@@ -98,39 +97,8 @@ EnumTypeTraversalImpl! {
 }
 
 impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Rc<T> {
-    fn try_fold_with<F: FallibleTypeFolder<I>>(mut self, folder: &mut F) -> Result<Self, F::Error> {
-        // 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 try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        self.try_map_id(|value| value.try_fold_with(folder))
     }
 }
 
@@ -141,39 +109,8 @@ impl<I: Interner, T: TypeVisitable<I>> TypeVisitable<I> for Rc<T> {
 }
 
 impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Arc<T> {
-    fn try_fold_with<F: FallibleTypeFolder<I>>(mut self, folder: &mut F) -> Result<Self, F::Error> {
-        // 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 try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        self.try_map_id(|value| value.try_fold_with(folder))
     }
 }