add variance support to TypeFoldable

Author: arielb1

src/librustc/ty/fold.rs

@@ -33,6 +33,7 @@
 
 use crate::hir::def_id::DefId;
 use crate::ty::{self, Binder, Ty, TyCtxt, TypeFlags};
+use crate::ty::subst::SubstsRef;
 
 use std::collections::BTreeMap;
 use std::fmt;
@@ -164,6 +165,35 @@ pub trait TypeFolder<'gcx: 'tcx, 'tcx> : Sized {
         t.super_fold_with(self)
     }
 
+    #[inline]
+    /// If `false` - the default - then `ty::Invariant` might be used instead of the
+    /// correct variance when folding an item with a variance.
+    ///
+    /// Otherwise, the correct variance is looked up from the tcx, which can
+    /// be a performance and cycle hazard.
+    fn use_variances(&self) -> bool {
+        false
+    }
+
+    #[inline]
+    fn fold_item_substs(&mut self, item_def_id: DefId, substs: SubstsRef<'tcx>)
+                        -> Result<SubstsRef<'tcx>, Self::Error>
+    {
+        if self.use_variances() {
+            let variances = self.tcx().variances_of(item_def_id);
+            ty::subst::fold_with_variances(self, &variances, substs)
+        } else {
+            substs.fold_with(self)
+        }
+    }
+
+    #[inline]
+    fn fold_with_variance<T>(&mut self, _variance: ty::Variance, t: &T) -> Result<T, Self::Error>
+        where T : TypeFoldable<'tcx>
+    {
+        t.fold_with(self)
+    }
+
     fn fold_ty(&mut self, t: Ty<'tcx>) -> Result<Ty<'tcx>, Self::Error> {
         t.super_fold_with(self)
     }

src/librustc/ty/structural_impls.rs

@@ -3,6 +3,7 @@
 //! hand, though we've recently added some macros (e.g.,
 //! `BraceStructLiftImpl!`) to help with the tedium.
 
+use crate::hir;
 use crate::mir::ProjectionKind;
 use crate::mir::interpret::ConstValue;
 use crate::ty::{self, Lift, Ty, TyCtxt};
@@ -12,6 +13,7 @@ use smallvec::SmallVec;
 use crate::mir::interpret;
 
 use std::rc::Rc;
+use std::iter;
 
 ///////////////////////////////////////////////////////////////////////////
 // Atomic structs
@@ -770,16 +772,24 @@ impl<'tcx> TypeFoldable<'tcx> for Ty<'tcx> {
             ty::RawPtr(tm) => ty::RawPtr(tm.fold_with(folder)?),
             ty::Array(typ, sz) => ty::Array(typ.fold_with(folder)?, sz.fold_with(folder)?),
             ty::Slice(typ) => ty::Slice(typ.fold_with(folder)?),
-            ty::Adt(tid, substs) => ty::Adt(tid, substs.fold_with(folder)?),
-            ty::Dynamic(ref trait_ty, ref region) =>
-                ty::Dynamic(trait_ty.fold_with(folder)?, region.fold_with(folder)?),
+            ty::Adt(tid, substs) => {
+                ty::Adt(tid, folder.fold_item_substs(tid.did, substs)?)
+            }
+            ty::Dynamic(ref trait_ty, ref region) => {
+                let principal = trait_ty.fold_with(folder)?;
+                let region_bound = folder.fold_with_variance(ty::Contravariant, region)?;
+                ty::Dynamic(principal, region_bound)
+            }
             ty::Tuple(ts) => ty::Tuple(ts.fold_with(folder)?),
             ty::FnDef(def_id, substs) => {
-                ty::FnDef(def_id, substs.fold_with(folder)?)
+                ty::FnDef(def_id, folder.fold_item_substs(def_id, substs)?)
             }
             ty::FnPtr(f) => ty::FnPtr(f.fold_with(folder)?),
             ty::Ref(ref r, ty, mutbl) => {
-                ty::Ref(r.fold_with(folder)?, ty.fold_with(folder)?, mutbl)
+                let r = folder.fold_with_variance(ty::Contravariant, r)?;
+                // Fold the type as a TypeAndMut to get the correct variance.
+                let mt = ty::TypeAndMut { ty, mutbl }.fold_with(folder)?;
+                ty::Ref(r, mt.ty, mt.mutbl)
             }
             ty::Generator(did, substs, movability) => {
                 ty::Generator(
@@ -864,9 +874,31 @@ impl<'tcx> TypeFoldable<'tcx> for Ty<'tcx> {
     }
 }
 
-BraceStructTypeFoldableImpl! {
-    impl<'tcx> TypeFoldable<'tcx> for ty::TypeAndMut<'tcx> {
-        ty, mutbl
+
+
+impl<'tcx> TypeFoldable<'tcx> for ty::TypeAndMut<'tcx> {
+    fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, folder: &mut F)
+                                                              -> Result<Self, F::Error>
+    {
+        let ty::TypeAndMut { ty, mutbl } = self;
+        let variance = match mutbl {
+            hir::Mutability::MutImmutable => ty::Covariant,
+            hir::Mutability::MutMutable => ty::Invariant,
+        };
+
+        Ok(ty::TypeAndMut {
+            ty: folder.fold_with_variance(variance, ty)?,
+            mutbl: mutbl.fold_with(folder)?,
+        })
+    }
+
+    fn super_visit_with<V: TypeVisitor<'tcx>>(&self, visitor: &mut V)
+                                              -> Result<(), V::Error>
+   {
+        let ty::TypeAndMut { ty, mutbl } = self;
+
+        ty.visit_with(visitor)?;
+        mutbl.visit_with(visitor)
     }
 }
 
@@ -876,9 +908,45 @@ BraceStructTypeFoldableImpl! {
     }
 }
 
-BraceStructTypeFoldableImpl! {
-    impl<'tcx> TypeFoldable<'tcx> for ty::FnSig<'tcx> {
-        inputs_and_output, c_variadic, unsafety, abi
+impl<'tcx> TypeFoldable<'tcx> for ty::FnSig<'tcx> {
+    fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, folder: &mut F)
+                                                              -> Result<Self, F::Error>
+    {
+        let ty::FnSig { inputs_and_output, c_variadic, unsafety, abi } = self;
+
+        let inputs_and_output = if folder.use_variances() {
+            let inputs_and_output = self.inputs().iter().cloned()
+                .map(|x| (x, false))
+                .chain(iter::once((self.output(), true)))
+                .map(|(a, is_output)| {
+                    if is_output {
+                        a.fold_with(folder)
+                    } else {
+                        folder.fold_with_variance(ty::Contravariant, &a)
+                    }
+                }).collect::<Result<SmallVec<[_; 8]>, _>>()?;
+            folder.tcx().intern_type_list(&inputs_and_output)
+        } else {
+            folder.fold_with_variance(ty::Invariant, inputs_and_output)?
+        };
+
+        Ok(ty::FnSig {
+            inputs_and_output,
+            c_variadic: *c_variadic,
+            unsafety: *unsafety,
+            abi: *abi,
+        })
+    }
+
+    fn super_visit_with<V: TypeVisitor<'tcx>>(&self, visitor: &mut V)
+                                              -> Result<(), V::Error>
+    {
+        let ty::FnSig { inputs_and_output, c_variadic, unsafety, abi } = self;
+
+        inputs_and_output.visit_with(visitor)?;
+        c_variadic.visit_with(visitor)?;
+        unsafety.visit_with(visitor)?;
+        abi.visit_with(visitor)
     }
 }
 

src/librustc/ty/subst.rs

@@ -329,8 +329,9 @@ impl<'tcx> TypeFoldable<'tcx> for SubstsRef<'tcx> {
     fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, folder: &mut F)
                                                               -> Result<Self, F::Error>
     {
-        let params = self.iter().map(|k| k.fold_with(folder))
-            .collect::<Result<SmallVec<[_; 8]>, _>>()?;
+        let params = self.iter().map(|k| {
+            folder.fold_with_variance(ty::Invariant, k)
+        }).collect::<Result<SmallVec<[_; 8]>, _>>()?;
 
         // If folding doesn't change the substs, it's faster to avoid
         // calling `mk_substs` and instead reuse the existing substs.
@@ -349,6 +350,27 @@ impl<'tcx> TypeFoldable<'tcx> for SubstsRef<'tcx> {
     }
 }
 
+pub fn fold_with_variances<'gcx, 'tcx, F: TypeFolder<'gcx, 'tcx>>(
+    folder: &mut F,
+    variances: &[ty::Variance],
+    substs: SubstsRef<'tcx>)
+    -> Result<SubstsRef<'tcx>, F::Error>
+{
+    assert_eq!(substs.len(), variances.len());
+
+    let params = substs.iter().zip(variances.iter()).map(|(k, v)| {
+        folder.fold_with_variance(*v, k)
+    }).collect::<Result<SmallVec<[_; 8]>, _>>()?;
+
+    // If folding doesn't change the substs, it's faster to avoid
+    // calling `mk_substs` and instead reuse the existing substs.
+    if params[..] == substs[..] {
+        Ok(substs)
+    } else {
+        Ok(folder.tcx().intern_substs(&params))
+    }
+}
+
 impl<'tcx> serialize::UseSpecializedDecodable for SubstsRef<'tcx> {}
 
 ///////////////////////////////////////////////////////////////////////////