Make TypeFolder::fold_* return Result

compiler/rustc_data_structures/src/functor.rs

@@ -2,12 +2,16 @@ use rustc_index::vec::{Idx, IndexVec};
 use std::mem;
 use std::ptr;
 
-pub trait IdFunctor {
+pub trait IdFunctor: Sized {
     type Inner;
 
     fn map_id<F>(self, f: F) -> Self
     where
         F: FnMut(Self::Inner) -> Self::Inner;
+
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
+    where
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>;
 }
 
 impl<T> IdFunctor for Box<T> {
@@ -31,6 +35,25 @@ impl<T> IdFunctor for Box<T> {
             raw.assume_init()
         }
     }
+
+    #[inline]
+    fn try_map_id<F, E>(self, mut f: F) -> Result<Self, E>
+    where
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
+    {
+        let raw = Box::into_raw(self);
+        Ok(unsafe {
+            // SAFETY: The raw pointer points to a valid value of type `T`.
+            let value = ptr::read(raw);
+            // SAFETY: Converts `Box<T>` to `Box<MaybeUninit<T>>` which is the
+            // inverse of `Box::assume_init()` and should be safe.
+            let mut raw: Box<mem::MaybeUninit<T>> = Box::from_raw(raw.cast());
+            // SAFETY: Write the mapped value back into the `Box`.
+            ptr::write(raw.as_mut_ptr(), f(value)?);
+            // SAFETY: We just initialized `raw`.
+            raw.assume_init()
+        })
+    }
 }
 
 impl<T> IdFunctor for Vec<T> {
@@ -55,6 +78,35 @@ impl<T> IdFunctor for Vec<T> {
         }
         self
     }
+
+    #[inline]
+    fn try_map_id<F, E>(mut self, mut f: F) -> Result<Self, E>
+    where
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
+    {
+        // FIXME: We don't really care about panics here and leak
+        // far more than we should, but that should be fine for now.
+        let len = self.len();
+        let mut error = Ok(());
+        unsafe {
+            self.set_len(0);
+            let start = self.as_mut_ptr();
+            for i in 0..len {
+                let p = start.add(i);
+                match f(ptr::read(p)) {
+                    Ok(value) => ptr::write(p, value),
+                    Err(err) => {
+                        error = Err(err);
+                        break;
+                    }
+                }
+            }
+            // Even if we encountered an error, set the len back
+            // so we don't leak memory.
+            self.set_len(len);
+        }
+        error.map(|()| self)
+    }
 }
 
 impl<T> IdFunctor for Box<[T]> {
@@ -67,6 +119,14 @@ impl<T> IdFunctor for Box<[T]> {
     {
         Vec::from(self).map_id(f).into()
     }
+
+    #[inline]
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
+    where
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
+    {
+        Vec::from(self).try_map_id(f).map(Into::into)
+    }
 }
 
 impl<I: Idx, T> IdFunctor for IndexVec<I, T> {
@@ -79,4 +139,12 @@ impl<I: Idx, T> IdFunctor for IndexVec<I, T> {
     {
         IndexVec::from_raw(self.raw.map_id(f))
     }
+
+    #[inline]
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
+    where
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
+    {
+        self.raw.try_map_id(f).map(IndexVec::from_raw)
+    }
 }

compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

@@ -293,7 +293,7 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
         self.tcx
     }
 
-    fn fold_binder<T>(&mut self, t: ty::Binder<'tcx, T>) -> ty::Binder<'tcx, T>
+    fn fold_binder<T>(&mut self, t: ty::Binder<'tcx, T>) -> Result<ty::Binder<'tcx, T>, Self::Error>
     where
         T: TypeFoldable<'tcx>,
     {
@@ -303,13 +303,13 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
         t
     }
 
-    fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
+    fn fold_region(&mut self, r: ty::Region<'tcx>) -> Result<ty::Region<'tcx>, Self::Error> {
         match *r {
             ty::ReLateBound(index, ..) => {
                 if index >= self.binder_index {
                     bug!("escaping late-bound region during canonicalization");
                 } else {
-                    r
+                    Ok(r)
                 }
             }
 
@@ -327,19 +327,19 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                     vid, r
                 );
                 let r = self.tcx.reuse_or_mk_region(r, ty::ReVar(resolved_vid));
-                self.canonicalize_region_mode.canonicalize_free_region(self, r)
+                Ok(self.canonicalize_region_mode.canonicalize_free_region(self, r))
             }
 
             ty::ReStatic
             | ty::ReEarlyBound(..)
             | ty::ReFree(_)
             | ty::ReEmpty(_)
             | ty::RePlaceholder(..)
-            | ty::ReErased => self.canonicalize_region_mode.canonicalize_free_region(self, r),
+            | ty::ReErased => Ok(self.canonicalize_region_mode.canonicalize_free_region(self, r)),
         }
     }
 
-    fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
+    fn fold_ty(&mut self, t: Ty<'tcx>) -> Result<Ty<'tcx>, Self::Error> {
         match *t.kind() {
             ty::Infer(ty::TyVar(vid)) => {
                 debug!("canonical: type var found with vid {:?}", vid);
@@ -355,40 +355,40 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                     Err(mut ui) => {
                         // FIXME: perf problem described in #55921.
                         ui = ty::UniverseIndex::ROOT;
-                        self.canonicalize_ty_var(
+                        Ok(self.canonicalize_ty_var(
                             CanonicalVarInfo {
                                 kind: CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui)),
                             },
                             t,
-                        )
+                        ))
                     }
                 }
             }
 
-            ty::Infer(ty::IntVar(_)) => self.canonicalize_ty_var(
+            ty::Infer(ty::IntVar(_)) => Ok(self.canonicalize_ty_var(
                 CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Int) },
                 t,
-            ),
+            )),
 
-            ty::Infer(ty::FloatVar(_)) => self.canonicalize_ty_var(
+            ty::Infer(ty::FloatVar(_)) => Ok(self.canonicalize_ty_var(
                 CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Float) },
                 t,
-            ),
+            )),
 
             ty::Infer(ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_)) => {
                 bug!("encountered a fresh type during canonicalization")
             }
 
-            ty::Placeholder(placeholder) => self.canonicalize_ty_var(
+            ty::Placeholder(placeholder) => Ok(self.canonicalize_ty_var(
                 CanonicalVarInfo { kind: CanonicalVarKind::PlaceholderTy(placeholder) },
                 t,
-            ),
+            )),
 
             ty::Bound(debruijn, _) => {
                 if debruijn >= self.binder_index {
                     bug!("escaping bound type during canonicalization")
                 } else {
-                    t
+                    Ok(t)
                 }
             }
 
@@ -419,13 +419,16 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                 if t.flags().intersects(self.needs_canonical_flags) {
                     t.super_fold_with(self)
                 } else {
-                    t
+                    Ok(t)
                 }
             }
         }
     }
 
-    fn fold_const(&mut self, ct: &'tcx ty::Const<'tcx>) -> &'tcx ty::Const<'tcx> {
+    fn fold_const(
+        &mut self,
+        ct: &'tcx ty::Const<'tcx>,
+    ) -> Result<&'tcx ty::Const<'tcx>, Self::Error> {
         match ct.val {
             ty::ConstKind::Infer(InferConst::Var(vid)) => {
                 debug!("canonical: const var found with vid {:?}", vid);
@@ -440,10 +443,10 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                     Err(mut ui) => {
                         // FIXME: perf problem described in #55921.
                         ui = ty::UniverseIndex::ROOT;
-                        return self.canonicalize_const_var(
+                        return Ok(self.canonicalize_const_var(
                             CanonicalVarInfo { kind: CanonicalVarKind::Const(ui) },
                             ct,
-                        );
+                        ));
                     }
                 }
             }
@@ -454,20 +457,20 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
                 if debruijn >= self.binder_index {
                     bug!("escaping bound type during canonicalization")
                 } else {
-                    return ct;
+                    return Ok(ct);
                 }
             }
             ty::ConstKind::Placeholder(placeholder) => {
-                return self.canonicalize_const_var(
+                return Ok(self.canonicalize_const_var(
                     CanonicalVarInfo { kind: CanonicalVarKind::PlaceholderConst(placeholder) },
                     ct,
-                );
+                ));
             }
             _ => {}
         }
 
         let flags = FlagComputation::for_const(ct);
-        if flags.intersects(self.needs_canonical_flags) { ct.super_fold_with(self) } else { ct }
+        if flags.intersects(self.needs_canonical_flags) { ct.super_fold_with(self) } else { Ok(ct) }
     }
 }
 
@@ -516,7 +519,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
             indices: FxHashMap::default(),
             binder_index: ty::INNERMOST,
         };
-        let out_value = value.fold_with(&mut canonicalizer);
+        let out_value = value.fold_with(&mut canonicalizer).into_ok();
 
         // Once we have canonicalized `out_value`, it should not
         // contain anything that ties it to this inference context
@@ -634,7 +637,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         let infcx = self.infcx.expect("encountered ty-var without infcx");
         let bound_to = infcx.shallow_resolve(ty_var);
         if bound_to != ty_var {
-            self.fold_ty(bound_to)
+            self.fold_ty(bound_to).into_ok()
         } else {
             let var = self.canonical_var(info, ty_var.into());
             self.tcx().mk_ty(ty::Bound(self.binder_index, var.into()))
@@ -653,12 +656,12 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         let infcx = self.infcx.expect("encountered const-var without infcx");
         let bound_to = infcx.shallow_resolve(const_var);
         if bound_to != const_var {
-            self.fold_const(bound_to)
+            self.fold_const(bound_to).into_ok()
         } else {
             let var = self.canonical_var(info, const_var.into());
             self.tcx().mk_const(ty::Const {
                 val: ty::ConstKind::Bound(self.binder_index, var),
-                ty: self.fold_ty(const_var.ty),
+                ty: self.fold_ty(const_var.ty).into_ok(),
             })
         }
     }