Only visit types once when walking the type tree

This fixes #72408.

Nested closures were resulting in exponential compilation time.

Author: VFLashM

src/librustc_infer/infer/outlives/verify.rs

@@ -2,6 +2,7 @@ use crate::infer::outlives::env::RegionBoundPairs;
 use crate::infer::{GenericKind, VerifyBound};
 use crate::traits;
 use rustc_data_structures::captures::Captures;
+use rustc_data_structures::fx::FxHashSet;
 use rustc_hir::def_id::DefId;
 use rustc_middle::ty::subst::{GenericArg, GenericArgKind, InternalSubsts, Subst};
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -32,16 +33,17 @@ 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 = FxHashSet::default();
         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 FxHashSet<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),
@@ -51,9 +53,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.
@@ -67,7 +69,7 @@ impl<'cx, 'tcx> VerifyBoundCx<'cx, 'tcx> {
                     ),
                 }
             }
-            _ => self.recursive_bound(ty.into()),
+            _ => self.recursive_bound(ty.into(), visited),
         }
     }
 
@@ -138,7 +140,7 @@ 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 FxHashSet<GenericArg<'tcx>>) -> VerifyBound<'tcx> {
         debug!("projection_bound(projection_ty={:?})", projection_ty);
 
         let projection_ty_as_ty =
@@ -167,21 +169,21 @@ 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 FxHashSet<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.

src/librustc_middle/ty/outlives.rs

@@ -4,6 +4,7 @@
 
 use crate::ty::subst::{GenericArg, GenericArgKind};
 use crate::ty::{self, Ty, TyCtxt, TypeFoldable};
+use rustc_data_structures::fx::FxHashSet;
 use smallvec::SmallVec;
 
 #[derive(Debug)]
@@ -50,12 +51,13 @@ 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 = FxHashSet::default();
+        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 FxHashSet<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 +75,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 +137,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 = FxHashSet::default();
+                    compute_components_recursive(tcx, ty.into(), &mut subcomponents, &mut subvisited);
                     out.push(Component::EscapingProjection(subcomponents.into_iter().collect()));
                 }
             }
@@ -177,7 +180,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 +189,12 @@ fn compute_components_recursive(
     tcx: TyCtxt<'tcx>,
     parent: GenericArg<'tcx>,
     out: &mut SmallVec<[Component<'tcx>; 4]>,
+    visited: &mut FxHashSet<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 +203,7 @@ fn compute_components_recursive(
                 }
             }
             GenericArgKind::Const(_) => {
-                compute_components_recursive(tcx, child, out);
+                compute_components_recursive(tcx, child, out, visited);
             }
         }
     }

src/librustc_middle/ty/walk.rs

@@ -4,6 +4,7 @@
 use crate::ty;
 use crate::ty::subst::{GenericArg, GenericArgKind};
 use smallvec::{self, SmallVec};
+use rustc_data_structures::fx::FxHashSet;
 
 // The TypeWalker's stack is hot enough that it's worth going to some effort to
 // avoid heap allocations.
@@ -12,11 +13,20 @@ type TypeWalkerStack<'tcx> = SmallVec<[GenericArg<'tcx>; 8]>;
 pub struct TypeWalker<'tcx> {
     stack: TypeWalkerStack<'tcx>,
     last_subtree: usize,
+    visited: FxHashSet<GenericArg<'tcx>>,
 }
 
+/// An iterator for walking the type tree.
+///
+/// It's very easy to produce a deeply
+/// nested type tree with a lot of
+/// identical subtrees. In order to work efficiently
+/// in this situation walker only visits each type once.
+/// It maintains a set of visited types and
+/// skips any types that are already there.
 impl<'tcx> TypeWalker<'tcx> {
-    pub fn new(root: GenericArg<'tcx>) -> TypeWalker<'tcx> {
-        TypeWalker { stack: smallvec![root], last_subtree: 1 }
+    pub fn new(root: GenericArg<'tcx>) -> Self {
+        Self { stack: smallvec![root], last_subtree: 1, visited: FxHashSet::default() }
     }
 
     /// Skips the subtree corresponding to the last type
@@ -41,11 +51,15 @@ impl<'tcx> Iterator for TypeWalker<'tcx> {
 
     fn next(&mut self) -> Option<GenericArg<'tcx>> {
         debug!("next(): stack={:?}", self.stack);
-        let next = self.stack.pop()?;
-        self.last_subtree = self.stack.len();
-        push_inner(&mut self.stack, next);
-        debug!("next: stack={:?}", self.stack);
-        Some(next)
+        loop {
+            let next = self.stack.pop()?;
+            self.last_subtree = self.stack.len();
+            if self.visited.insert(next) {
+                push_inner(&mut self.stack, next);
+                debug!("next: stack={:?}", self.stack);
+                return Some(next);
+            }
+        }
     }
 }
 
@@ -67,9 +81,16 @@ impl GenericArg<'tcx> {
     /// Iterator that walks the immediate children of `self`. Hence
     /// `Foo<Bar<i32>, u32>` yields the sequence `[Bar<i32>, u32]`
     /// (but not `i32`, like `walk`).
-    pub fn walk_shallow(self) -> impl Iterator<Item = GenericArg<'tcx>> {
+    ///
+    /// Iterator only walks items once.
+    /// It accepts visited set, updates it with all visited types
+    /// and skips any types that are already there.
+    pub fn walk_shallow(self, visited: &mut FxHashSet<GenericArg<'tcx>>) -> impl Iterator<Item = GenericArg<'tcx>> {
         let mut stack = SmallVec::new();
         push_inner(&mut stack, self);
+        stack.retain(|a| {
+            visited.insert(*a)
+        });
         stack.into_iter()
     }
 }

src/test/ui/closures/issue-72408-nested-closures-exponential.rs

@@ -0,0 +1,59 @@
+// build-pass
+
+// Closures include captured types twice in a type tree.
+//
+// Wrapping one closure with another leads to doubling
+// the amount of types in the type tree.
+//
+// This test ensures that rust can handle
+// deeply nested type trees with a lot
+// of duplicated subtrees.
+
+fn dup(f: impl Fn(i32) -> i32) -> impl Fn(i32) -> i32 {
+    move |a| f(a * 2)
+}
+
+fn main() {
+    let f = |a| a;
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    // Compiler dies around here if it tries
+    // to walk the tree exhaustively.
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+    let f = dup(f);
+
+    println!("Type size was at least {}", f(1));
+}