1 files changed, 103 insertions, 0 deletions
diff --git a/compiler/rustc_type_ir/src/fold.rs b/compiler/rustc_type_ir/src/fold.rs
index e7a6831f5..fc56400df 100644
--- a/compiler/rustc_type_ir/src/fold.rs
+++ b/compiler/rustc_type_ir/src/fold.rs
@@ -44,6 +44,11 @@
 //!     - ty.super_fold_with(folder)
 //! - u.fold_with(folder)
 //! ```
+
+use rustc_data_structures::sync::Lrc;
+use rustc_index::{Idx, IndexVec};
+use std::mem;
+
 use crate::{visit::TypeVisitable, Interner};
 
 /// This trait is implemented for every type that can be folded,
@@ -242,3 +247,101 @@ where
         Ok(self.fold_predicate(p))
     }
 }
+
+///////////////////////////////////////////////////////////////////////////
+// Traversal implementations.
+
+impl<I: Interner, T: TypeFoldable<I>, U: TypeFoldable<I>> TypeFoldable<I> for (T, U) {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<(T, U), F::Error> {
+        Ok((self.0.try_fold_with(folder)?, self.1.try_fold_with(folder)?))
+    }
+}
+
+impl<I: Interner, A: TypeFoldable<I>, B: TypeFoldable<I>, C: TypeFoldable<I>> TypeFoldable<I>
+    for (A, B, C)
+{
+    fn try_fold_with<F: FallibleTypeFolder<I>>(
+        self,
+        folder: &mut F,
+    ) -> Result<(A, B, C), F::Error> {
+        Ok((
+            self.0.try_fold_with(folder)?,
+            self.1.try_fold_with(folder)?,
+            self.2.try_fold_with(folder)?,
+        ))
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Option<T> {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        Ok(match self {
+            Some(v) => Some(v.try_fold_with(folder)?),
+            None => None,
+        })
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>, E: TypeFoldable<I>> TypeFoldable<I> for Result<T, E> {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        Ok(match self {
+            Ok(v) => Ok(v.try_fold_with(folder)?),
+            Err(e) => Err(e.try_fold_with(folder)?),
+        })
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Lrc<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 `Lrc` or indeed clone
+        // the contained type.
+        unsafe {
+            // First step is to ensure that we have a unique reference to
+            // the contained type, which `Lrc::make_mut` will accomplish (by
+            // allocating a new `Lrc` and cloning the `T` only if required).
+            // This is done *before* casting to `Lrc<ManuallyDrop<T>>` so that
+            // panicking during `make_mut` does not leak the `T`.
+            Lrc::make_mut(&mut self);
+
+            // Casting to `Lrc<ManuallyDrop<T>>` is safe because `ManuallyDrop`
+            // is `repr(transparent)`.
+            let ptr = Lrc::into_raw(self).cast::<mem::ManuallyDrop<T>>();
+            let mut unique = Lrc::from_raw(ptr);
+
+            // Call to `Lrc::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 = Lrc::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 = owned.try_fold_with(folder)?;
+            *slot = mem::ManuallyDrop::new(folded);
+
+            // Cast back to `Lrc<T>`.
+            Ok(Lrc::from_raw(Lrc::into_raw(unique).cast()))
+        }
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Box<T> {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(mut self, folder: &mut F) -> Result<Self, F::Error> {
+        *self = (*self).try_fold_with(folder)?;
+        Ok(self)
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Vec<T> {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        self.into_iter().map(|t| t.try_fold_with(folder)).collect()
+    }
+}
+
+impl<I: Interner, T: TypeFoldable<I>, Ix: Idx> TypeFoldable<I> for IndexVec<Ix, T> {
+    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
+        self.raw.try_fold_with(folder).map(IndexVec::from_raw)
+    }
+}