Auto merge of #91230 - eggyal:fallible-type-fold, r=jackh726

Make `TypeFolder::fold_*` return `Result`

Implements rust-lang/compiler-team#432.

Initially this is just a rebase of `@LeSeulArtichaut's` work in #85469 (abandoned; see #85485 (comment)).  At that time, it caused a regression in performance that required some further exploration... with this rebased PR bors can hopefully report some perf analysis from which we can investigate further (if the regression is indeed still present).

r? `@jackh726` cc `@nikomatsakis`

compiler/rustc_const_eval/src/lib.rs

@@ -23,6 +23,7 @@ Rust MIR: a lowered representation of Rust.
 #![feature(trusted_len)]
 #![feature(trusted_step)]
 #![feature(try_blocks)]
+#![feature(unwrap_infallible)]
 #![recursion_limit = "256"]
 
 #[macro_use]

compiler/rustc_const_eval/src/transform/validate.rs

@@ -95,7 +95,8 @@ pub fn equal_up_to_regions(
                 // Leave consts and types unchanged.
                 ct_op: |ct| ct,
                 ty_op: |ty| ty,
-            }),
+            })
+            .into_ok(),
         )
     };
     tcx.infer_ctxt().enter(|infcx| infcx.can_eq(param_env, normalize(src), normalize(dest)).is_ok())

compiler/rustc_data_structures/src/functor.rs

@@ -1,45 +1,52 @@
 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
+    #[inline]
+    fn map_id<F>(self, mut f: F) -> Self
+    where
+        F: FnMut(Self::Inner) -> Self::Inner,
+    {
+        self.try_map_id::<_, !>(|value| Ok(f(value))).into_ok()
+    }
+
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
     where
-        F: FnMut(Self::Inner) -> Self::Inner;
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>;
 }
 
 impl<T> IdFunctor for Box<T> {
     type Inner = T;
 
     #[inline]
-    fn map_id<F>(self, mut f: F) -> Self
+    fn try_map_id<F, E>(self, mut f: F) -> Result<Self, E>
     where
-        F: FnMut(Self::Inner) -> Self::Inner,
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
     {
         let raw = Box::into_raw(self);
-        unsafe {
+        Ok(unsafe {
             // SAFETY: The raw pointer points to a valid value of type `T`.
-            let value = ptr::read(raw);
+            let value = raw.read();
             // 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`.
-            raw.write(f(value));
+            raw.write(f(value)?);
             // SAFETY: We just initialized `raw`.
             raw.assume_init()
-        }
+        })
     }
 }
 
 impl<T> IdFunctor for Vec<T> {
     type Inner = T;
 
     #[inline]
-    fn map_id<F>(mut self, mut f: F) -> Self
+    fn try_map_id<F, E>(mut self, mut f: F) -> Result<Self, E>
     where
-        F: FnMut(Self::Inner) -> Self::Inner,
+        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.
@@ -49,34 +56,49 @@ impl<T> IdFunctor for Vec<T> {
             let start = self.as_mut_ptr();
             for i in 0..len {
                 let p = start.add(i);
-                ptr::write(p, f(ptr::read(p)));
+                match f(p.read()) {
+                    Ok(val) => p.write(val),
+                    Err(err) => {
+                        // drop all other elements in self
+                        // (current element was "moved" into the call to f)
+                        for j in (0..i).chain(i + 1..len) {
+                            start.add(j).drop_in_place();
+                        }
+
+                        // returning will drop self, releasing the allocation
+                        // (len is 0 so elements will not be re-dropped)
+                        return Err(err);
+                    }
+                }
             }
+            // Even if we encountered an error, set the len back
+            // so we don't leak memory.
             self.set_len(len);
         }
-        self
+        Ok(self)
     }
 }
 
 impl<T> IdFunctor for Box<[T]> {
     type Inner = T;
 
     #[inline]
-    fn map_id<F>(self, f: F) -> Self
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
     where
-        F: FnMut(Self::Inner) -> Self::Inner,
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
     {
-        Vec::from(self).map_id(f).into()
+        Vec::from(self).try_map_id(f).map(Into::into)
     }
 }
 
 impl<I: Idx, T> IdFunctor for IndexVec<I, T> {
     type Inner = T;
 
     #[inline]
-    fn map_id<F>(self, f: F) -> Self
+    fn try_map_id<F, E>(self, f: F) -> Result<Self, E>
     where
-        F: FnMut(Self::Inner) -> Self::Inner,
+        F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
     {
-        IndexVec::from_raw(self.raw.map_id(f))
+        self.raw.try_map_id(f).map(IndexVec::from_raw)
     }
 }

compiler/rustc_data_structures/src/lib.rs

@@ -25,6 +25,7 @@
 #![feature(once_cell)]
 #![feature(test)]
 #![feature(thread_id_value)]
+#![feature(unwrap_infallible)]
 #![allow(rustc::default_hash_types)]
 #![deny(unaligned_references)]
 

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

@@ -278,7 +278,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>,
     {
@@ -288,13 +288,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)
                 }
             }
 
@@ -311,19 +311,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);
@@ -339,40 +339,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)
                 }
             }
 
@@ -403,13 +403,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);
@@ -424,10 +427,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,
-                        );
+                        ));
                     }
                 }
             }
@@ -438,20 +441,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) }
     }
 }
 
@@ -500,7 +503,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
@@ -618,7 +621,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         let infcx = self.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()))
@@ -637,12 +640,12 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         let infcx = self.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(),
             })
         }
     }