Auto merge of rust-lang#72412 - VFLashM:issue-72408-nested-closures-exponential, r=tmandry

Issue 72408 nested closures exponential

This fixes rust-lang#72408.

Nested closures were resulting in exponential compilation time.

This PR is enhancing asymptotic complexity, but also increasing the constant, so I would love to see perf run results.

Author: bors

Cargo.lock

@@ -3605,6 +3605,7 @@ dependencies = [
 name = "rustc_infer"
 version = "0.0.0"
 dependencies = [
+ "arrayvec",
  "rustc_ast",
  "rustc_data_structures",
  "rustc_errors",
@@ -3744,6 +3745,7 @@ dependencies = [
 name = "rustc_middle"
 version = "0.0.0"
 dependencies = [
+ "arrayvec",
  "bitflags",
  "chalk-ir",
  "measureme",

compiler/rustc_infer/Cargo.toml

@@ -21,4 +21,5 @@ rustc_serialize = { path = "../rustc_serialize" }
 rustc_span = { path = "../rustc_span" }
 rustc_target = { path = "../rustc_target" }
 smallvec = { version = "1.0", features = ["union", "may_dangle"] }
+arrayvec = { version = "0.5.1", default-features = false }
 rustc_ast = { path = "../rustc_ast" }

compiler/rustc_infer/src/infer/combine.rs

@@ -31,6 +31,9 @@ use super::unify_key::replace_if_possible;
 use super::unify_key::{ConstVarValue, ConstVariableValue};
 use super::unify_key::{ConstVariableOrigin, ConstVariableOriginKind};
 use super::{InferCtxt, MiscVariable, TypeTrace};
+use arrayvec::ArrayVec;
+use rustc_data_structures::fx::FxHashMap;
+use std::hash::Hash;
 
 use crate::traits::{Obligation, PredicateObligations};
 
@@ -44,6 +47,63 @@ use rustc_middle::ty::{self, InferConst, ToPredicate, Ty, TyCtxt, TypeFoldable};
 use rustc_middle::ty::{IntType, UintType};
 use rustc_span::DUMMY_SP;
 
+/// Small-storage-optimized implementation of a map
+/// made specifically for caching results.
+///
+/// Stores elements in a small array up to a certain length
+/// and switches to `HashMap` when that length is exceeded.
+enum MiniMap<K, V> {
+    Array(ArrayVec<[(K, V); 8]>),
+    Map(FxHashMap<K, V>),
+}
+
+impl<K: Eq + Hash, V> MiniMap<K, V> {
+    /// Creates an empty `MiniMap`.
+    pub fn new() -> Self {
+        MiniMap::Array(ArrayVec::new())
+    }
+
+    /// Inserts or updates value in the map.
+    pub fn insert(&mut self, key: K, value: V) {
+        match self {
+            MiniMap::Array(array) => {
+                for pair in array.iter_mut() {
+                    if pair.0 == key {
+                        pair.1 = value;
+                        return;
+                    }
+                }
+                if let Err(error) = array.try_push((key, value)) {
+                    let mut map: FxHashMap<K, V> = array.drain(..).collect();
+                    let (key, value) = error.element();
+                    map.insert(key, value);
+                    *self = MiniMap::Map(map);
+                }
+            }
+            MiniMap::Map(map) => {
+                map.insert(key, value);
+            }
+        }
+    }
+
+    /// Return value by key if any.
+    pub fn get(&self, key: &K) -> Option<&V> {
+        match self {
+            MiniMap::Array(array) => {
+                for pair in array {
+                    if pair.0 == *key {
+                        return Some(&pair.1);
+                    }
+                }
+                return None;
+            }
+            MiniMap::Map(map) => {
+                return map.get(key);
+            }
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct CombineFields<'infcx, 'tcx> {
     pub infcx: &'infcx InferCtxt<'infcx, 'tcx>,
@@ -379,6 +439,7 @@ impl<'infcx, 'tcx> CombineFields<'infcx, 'tcx> {
             needs_wf: false,
             root_ty: ty,
             param_env: self.param_env,
+            cache: MiniMap::new(),
         };
 
         let ty = match generalize.relate(ty, ty) {
@@ -438,6 +499,8 @@ struct Generalizer<'cx, 'tcx> {
     root_ty: Ty<'tcx>,
 
     param_env: ty::ParamEnv<'tcx>,
+
+    cache: MiniMap<Ty<'tcx>, RelateResult<'tcx, Ty<'tcx>>>,
 }
 
 /// Result from a generalization operation. This includes
@@ -535,13 +598,16 @@ impl TypeRelation<'tcx> for Generalizer<'_, 'tcx> {
     fn tys(&mut self, t: Ty<'tcx>, t2: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
         assert_eq!(t, t2); // we are abusing TypeRelation here; both LHS and RHS ought to be ==
 
+        if let Some(result) = self.cache.get(&t) {
+            return result.clone();
+        }
         debug!("generalize: t={:?}", t);
 
         // Check to see whether the type we are generalizing references
         // any other type variable related to `vid` via
         // subtyping. This is basically our "occurs check", preventing
         // us from creating infinitely sized types.
-        match *t.kind() {
+        let result = match *t.kind() {
             ty::Infer(ty::TyVar(vid)) => {
                 let vid = self.infcx.inner.borrow_mut().type_variables().root_var(vid);
                 let sub_vid = self.infcx.inner.borrow_mut().type_variables().sub_root_var(vid);
@@ -598,7 +664,10 @@ impl TypeRelation<'tcx> for Generalizer<'_, 'tcx> {
                 Ok(t)
             }
             _ => relate::super_relate_tys(self, t, t),
-        }
+        };
+
+        self.cache.insert(t, result.clone());
+        return result;
     }
 
     fn regions(

compiler/rustc_infer/src/infer/outlives/verify.rs

@@ -3,6 +3,7 @@ use crate::infer::{GenericKind, VerifyBound};
 use rustc_data_structures::captures::Captures;
 use rustc_hir::def_id::DefId;
 use rustc_middle::ty::subst::{GenericArg, GenericArgKind, Subst};
+use rustc_middle::ty::walk::MiniSet;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 
 /// The `TypeOutlives` struct has the job of "lowering" a `T: 'a`
@@ -31,16 +32,23 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
     /// Returns a "verify bound" that encodes what we know about
     /// `generic` and the regions it outlives.
     pub fn generic_bound(&self, generic: GenericKind<'tcx>) -> VerifyBound<'tcx> {
+        let mut visited = MiniSet::new();
         match generic {
             GenericKind::Param(param_ty) => self.param_bound(param_ty),
-            GenericKind::Projection(projection_ty) => self.projection_bound(projection_ty),
+            GenericKind::Projection(projection_ty) => {
+                self.projection_bound(projection_ty, &mut visited)
+            }
         }
     }
 
-    fn type_bound(&self, ty: Ty<'tcx>) -> VerifyBound<'tcx> {
+    fn type_bound(
+        &self,
+        ty: Ty<'tcx>,
+        visited: &mut MiniSet<GenericArg<'tcx>>,
+    ) -> VerifyBound<'tcx> {
         match *ty.kind() {
             ty::Param(p) => self.param_bound(p),
-            ty::Projection(data) => self.projection_bound(data),
+            ty::Projection(data) => self.projection_bound(data, visited),
             ty::FnDef(_, substs) => {
                 // HACK(eddyb) ignore lifetimes found shallowly in `substs`.
                 // This is inconsistent with `ty::Adt` (including all substs),
@@ -50,9 +58,9 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
                 let mut bounds = substs
                     .iter()
                     .filter_map(|child| match child.unpack() {
-                        GenericArgKind::Type(ty) => Some(self.type_bound(ty)),
+                        GenericArgKind::Type(ty) => Some(self.type_bound(ty, visited)),
                         GenericArgKind::Lifetime(_) => None,
-                        GenericArgKind::Const(_) => Some(self.recursive_bound(child)),
+                        GenericArgKind::Const(_) => Some(self.recursive_bound(child, visited)),
                     })
                     .filter(|bound| {
                         // Remove bounds that must hold, since they are not interesting.
@@ -66,7 +74,7 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
                     ),
                 }
             }
-            _ => self.recursive_bound(ty.into()),
+            _ => self.recursive_bound(ty.into(), visited),
         }
     }
 
@@ -137,7 +145,11 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
         self.declared_projection_bounds_from_trait(projection_ty)
     }
 
-    pub fn projection_bound(&self, projection_ty: ty::ProjectionTy<'tcx>) -> VerifyBound<'tcx> {
+    pub fn projection_bound(
+        &self,
+        projection_ty: ty::ProjectionTy<'tcx>,
+        visited: &mut MiniSet<GenericArg<'tcx>>,
+    ) -> VerifyBound<'tcx> {
         debug!("projection_bound(projection_ty={:?})", projection_ty);
 
         let projection_ty_as_ty =
@@ -166,21 +178,25 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
 
         // see the extensive comment in projection_must_outlive
         let ty = self.tcx.mk_projection(projection_ty.item_def_id, projection_ty.substs);
-        let recursive_bound = self.recursive_bound(ty.into());
+        let recursive_bound = self.recursive_bound(ty.into(), visited);
 
         VerifyBound::AnyBound(env_bounds.chain(trait_bounds).collect()).or(recursive_bound)
     }
 
-    fn recursive_bound(&self, parent: GenericArg<'tcx>) -> VerifyBound<'tcx> {
+    fn recursive_bound(
+        &self,
+        parent: GenericArg<'tcx>,
+        visited: &mut MiniSet<GenericArg<'tcx>>,
+    ) -> VerifyBound<'tcx> {
         let mut bounds = parent
-            .walk_shallow()
+            .walk_shallow(visited)
             .filter_map(|child| match child.unpack() {
-                GenericArgKind::Type(ty) => Some(self.type_bound(ty)),
+                GenericArgKind::Type(ty) => Some(self.type_bound(ty, visited)),
                 GenericArgKind::Lifetime(lt) => {
                     // Ignore late-bound regions.
                     if !lt.is_late_bound() { Some(VerifyBound::OutlivedBy(lt)) } else { None }
                 }
-                GenericArgKind::Const(_) => Some(self.recursive_bound(child)),
+                GenericArgKind::Const(_) => Some(self.recursive_bound(child, visited)),
             })
             .filter(|bound| {
                 // Remove bounds that must hold, since they are not interesting.

compiler/rustc_middle/Cargo.toml

@@ -28,5 +28,6 @@ rustc_ast = { path = "../rustc_ast" }
 rustc_span = { path = "../rustc_span" }
 chalk-ir = "0.21.0"
 smallvec = { version = "1.0", features = ["union", "may_dangle"] }
+arrayvec = { version = "0.5.1", default-features = false }
 measureme = "0.7.1"
 rustc_session = { path = "../rustc_session" }

compiler/rustc_middle/src/ty/outlives.rs

@@ -3,6 +3,7 @@
 // RFC for reference.
 
 use crate::ty::subst::{GenericArg, GenericArgKind};
+use crate::ty::walk::MiniSet;
 use crate::ty::{self, Ty, TyCtxt, TypeFoldable};
 use smallvec::SmallVec;
 
@@ -50,12 +51,18 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Push onto `out` all the things that must outlive `'a` for the condition
     /// `ty0: 'a` to hold. Note that `ty0` must be a **fully resolved type**.
     pub fn push_outlives_components(self, ty0: Ty<'tcx>, out: &mut SmallVec<[Component<'tcx>; 4]>) {
-        compute_components(self, ty0, out);
+        let mut visited = MiniSet::new();
+        compute_components(self, ty0, out, &mut visited);
         debug!("components({:?}) = {:?}", ty0, out);
     }
 }
 
-fn compute_components(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, out: &mut SmallVec<[Component<'tcx>; 4]>) {
+fn compute_components(
+    tcx: TyCtxt<'tcx>,
+    ty: Ty<'tcx>,
+    out: &mut SmallVec<[Component<'tcx>; 4]>,
+    visited: &mut MiniSet<GenericArg<'tcx>>,
+) {
     // Descend through the types, looking for the various "base"
     // components and collecting them into `out`. This is not written
     // with `collect()` because of the need to sometimes skip subtrees
@@ -73,31 +80,31 @@ fn compute_components(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, out: &mut SmallVec<[Compo
                 for child in substs {
                     match child.unpack() {
                         GenericArgKind::Type(ty) => {
-                            compute_components(tcx, ty, out);
+                            compute_components(tcx, ty, out, visited);
                         }
                         GenericArgKind::Lifetime(_) => {}
                         GenericArgKind::Const(_) => {
-                            compute_components_recursive(tcx, child, out);
+                            compute_components_recursive(tcx, child, out, visited);
                         }
                     }
                 }
             }
 
             ty::Array(element, _) => {
                 // Don't look into the len const as it doesn't affect regions
-                compute_components(tcx, element, out);
+                compute_components(tcx, element, out, visited);
             }
 
             ty::Closure(_, ref substs) => {
                 for upvar_ty in substs.as_closure().upvar_tys() {
-                    compute_components(tcx, upvar_ty, out);
+                    compute_components(tcx, upvar_ty, out, visited);
                 }
             }
 
             ty::Generator(_, ref substs, _) => {
                 // Same as the closure case
                 for upvar_ty in substs.as_generator().upvar_tys() {
-                    compute_components(tcx, upvar_ty, out);
+                    compute_components(tcx, upvar_ty, out, visited);
                 }
 
                 // We ignore regions in the generator interior as we don't
@@ -135,7 +142,8 @@ fn compute_components(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, out: &mut SmallVec<[Compo
                     // OutlivesProjectionComponents.  Continue walking
                     // through and constrain Pi.
                     let mut subcomponents = smallvec![];
-                    compute_components_recursive(tcx, ty.into(), &mut subcomponents);
+                    let mut subvisited = MiniSet::new();
+                    compute_components_recursive(tcx, ty.into(), &mut subcomponents, &mut subvisited);
                     out.push(Component::EscapingProjection(subcomponents.into_iter().collect()));
                 }
             }
@@ -177,7 +185,7 @@ fn compute_components(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, out: &mut SmallVec<[Compo
                 // the "bound regions list".  In our representation, no such
                 // list is maintained explicitly, because bound regions
                 // themselves can be readily identified.
-                compute_components_recursive(tcx, ty.into(), out);
+                compute_components_recursive(tcx, ty.into(), out, visited);
             }
         }
 }
@@ -186,11 +194,12 @@ fn compute_components_recursive(
     tcx: TyCtxt<'tcx>,
     parent: GenericArg<'tcx>,
     out: &mut SmallVec<[Component<'tcx>; 4]>,
+    visited: &mut MiniSet<GenericArg<'tcx>>,
 ) {
-    for child in parent.walk_shallow() {
+    for child in parent.walk_shallow(visited) {
         match child.unpack() {
             GenericArgKind::Type(ty) => {
-                compute_components(tcx, ty, out);
+                compute_components(tcx, ty, out, visited);
             }
             GenericArgKind::Lifetime(lt) => {
                 // Ignore late-bound regions.
@@ -199,7 +208,7 @@ fn compute_components_recursive(
                 }
             }
             GenericArgKind::Const(_) => {
-                compute_components_recursive(tcx, child, out);
+                compute_components_recursive(tcx, child, out, visited);
             }
         }
     }