Track closure signatures & kinds in freshened types

This allows caching closure signatures and kinds in the normal selection
and evaluation caches, and fixes the exponential worst-case in
@remram44's example, which is a part of #43787.

This improvement is complenentary to #43999 - they fix different cases.

2 files changed, 155 insertions, 60 deletions

diff --git a/src/librustc/infer/freshen.rs b/src/librustc/infer/freshen.rs
index b8b5a55f578..c274f8bda9f 100644
--- a/src/librustc/infer/freshen.rs
+++ b/src/librustc/infer/freshen.rs
@@ -19,10 +19,21 @@
 //! fact an unbound type variable, we want the match to be regarded as ambiguous, because depending
 //! on what type that type variable is ultimately assigned, the match may or may not succeed.
 //!
+//! To handle closures, freshened types also have to contain the signature and kind of any
+//! closure in the local inference context, as otherwise the cache key might be invalidated.
+//! The way this is done is somewhat hacky - the closure signature is appended to the substs,
+//! as well as the closure kind "encoded" as a type. Also, special handling is needed when
+//! the closure signature contains a reference to the original closure.
+//!
 //! Note that you should be careful not to allow the output of freshening to leak to the user in
 //! error messages or in any other form. Freshening is only really useful as an internal detail.
 //!
-//! __An important detail concerning regions.__ The freshener also replaces *all* regions with
+//! Because of the manipulation required to handle closures, doing arbitrary operations on
+//! freshened types is not recommended. However, in addition to doing equality/hash
+//! comparisons (for caching), it is possible to do a `ty::_match` operation between
+//! 2 freshened types - this works even with the closure encoding.
+//!
+//! __An important detail concerning regions.__ The freshener also replaces *all* free regions with
 //! 'erased. The reason behind this is that, in general, we do not take region relationships into
 //! account when making type-overloaded decisions. This is important because of the design of the
 //! region inferencer, which is not based on unification but rather on accumulating and then
@@ -32,7 +43,10 @@
 
 use ty::{self, Ty, TyCtxt, TypeFoldable};
 use ty::fold::TypeFolder;
+use ty::subst::Substs;
 use util::nodemap::FxHashMap;
+use hir::def_id::DefId;
+
 use std::collections::hash_map::Entry;
 
 use super::InferCtxt;
@@ -42,6 +56,7 @@ pub struct TypeFreshener<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
     infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
     freshen_count: u32,
     freshen_map: FxHashMap<ty::InferTy, Ty<'tcx>>,
+    closure_set: Vec<DefId>,
 }
 
 impl<'a, 'gcx, 'tcx> TypeFreshener<'a, 'gcx, 'tcx> {
@@ -51,6 +66,7 @@ impl<'a, 'gcx, 'tcx> TypeFreshener<'a, 'gcx, 'tcx> {
             infcx,
             freshen_count: 0,
             freshen_map: FxHashMap(),
+            closure_set: vec![],
         }
     }
 
@@ -76,6 +92,88 @@ impl<'a, 'gcx, 'tcx> TypeFreshener<'a, 'gcx, 'tcx> {
             }
         }
     }
+
+    fn next_fresh<F>(&mut self,
+                     freshener: F)
+                     -> Ty<'tcx>
+        where F: FnOnce(u32) -> ty::InferTy,
+    {
+        let index = self.freshen_count;
+        self.freshen_count += 1;
+        self.infcx.tcx.mk_infer(freshener(index))
+    }
+
+    fn freshen_closure_like<M, C>(&mut self,
+                                  def_id: DefId,
+                                  substs: ty::ClosureSubsts<'tcx>,
+                                  t: Ty<'tcx>,
+                                  markers: M,
+                                  combine: C)
+                                  -> Ty<'tcx>
+        where M: FnOnce(&mut Self) -> (Ty<'tcx>, Ty<'tcx>),
+              C: FnOnce(&'tcx Substs<'tcx>) -> Ty<'tcx>
+    {
+        let tcx = self.infcx.tcx;
+
+        let closure_in_progress = self.infcx.in_progress_tables.map_or(false, |tables| {
+            tcx.hir.as_local_node_id(def_id).map_or(false, |closure_id| {
+                tables.borrow().local_id_root ==
+                    Some(DefId::local(tcx.hir.node_to_hir_id(closure_id).owner))
+            })
+        });
+
+        if !closure_in_progress {
+            // If this closure belongs to another infcx, its kind etc. were
+            // fully inferred and its signature/kind are exactly what's listed
+            // in its infcx. So we don't need to add the markers for them.
+            return t.super_fold_with(self);
+        }
+
+        // We are encoding a closure in progress. Because we want our freshening
+        // key to contain all inference information needed to make sense of our
+        // value, we need to encode the closure signature and kind. The way
+        // we do that is to add them as 2 variables to the closure substs,
+        // basically because it's there (and nobody cares about adding extra stuff
+        // to substs).
+        //
+        // This means the "freshened" closure substs ends up looking like
+        //     fresh_substs = [PARENT_SUBSTS* ; UPVARS* ; SIG_MARKER ; KIND_MARKER]
+        let (marker_1, marker_2) = if self.closure_set.contains(&def_id) {
+            // We found the closure def-id within its own signature. Just
+            // leave a new freshened type - any matching operations would
+            // have found and compared the exterior closure already to
+            // get here.
+            //
+            // In that case, we already know what the signature would
+            // be - the parent closure on the stack already contains a
+            // "copy" of the signature, so there is no reason to encode
+            // it again for injectivity. Just use a fresh type variable
+            // to make everything comparable.
+            //
+            // For example (closure kinds omitted for clarity)
+            //     t=[closure FOO sig=[closure BAR sig=[closure FOO ..]]]
+            // Would get encoded to
+            //     t=[closure FOO sig=[closure BAR sig=[closure FOO sig=$0]]]
+            //
+            // and we can decode by having
+            //     $0=[closure BAR {sig doesn't exist in decode}]
+            // and get
+            //     t=[closure FOO]
+            //     sig[FOO] = [closure BAR]
+            //     sig[BAR] = [closure FOO]
+            (self.next_fresh(ty::FreshTy), self.next_fresh(ty::FreshTy))
+        } else {
+            self.closure_set.push(def_id);
+            let markers = markers(self);
+            self.closure_set.pop();
+            markers
+        };
+
+        combine(tcx.mk_substs(
+            substs.substs.iter().map(|k| k.fold_with(self)).chain(
+                [marker_1, marker_2].iter().cloned().map(From::from)
+                    )))
+    }
 }
 
 impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
@@ -105,7 +203,8 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
     }
 
     fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
-        if !t.needs_infer() && !t.has_erasable_regions() {
+        if !t.needs_infer() && !t.has_erasable_regions() &&
+            !(t.has_closure_types() && self.infcx.in_progress_tables.is_some()) {
             return t;
         }
 
@@ -150,6 +249,51 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
                 t
             }
 
+            ty::TyClosure(def_id, substs) => {
+                self.freshen_closure_like(
+                    def_id, substs, t,
+                    |this| {
+                        // HACK: use a "random" integer type to mark the kind. Because
+                        // different closure kinds shouldn't get unified during
+                        // selection, the "subtyping" relationship (where any kind is
+                        // better than no kind) shouldn't  matter here, just that the
+                        // types are different.
+                        let closure_kind = this.infcx.closure_kind(def_id);
+                        let closure_kind_marker = match closure_kind {
+                            None => tcx.types.i8,
+                            Some(ty::ClosureKind::Fn) => tcx.types.i16,
+                            Some(ty::ClosureKind::FnMut) => tcx.types.i32,
+                            Some(ty::ClosureKind::FnOnce) => tcx.types.i64,
+                        };
+
+                        let closure_sig = this.infcx.fn_sig(def_id);
+                        (tcx.mk_fn_ptr(closure_sig.fold_with(this)),
+                         closure_kind_marker)
+                    },
+                    |substs| tcx.mk_closure(def_id, substs)
+                )
+            }
+
+            ty::TyGenerator(def_id, substs, interior) => {
+                self.freshen_closure_like(
+                    def_id, substs, t,
+                    |this| {
+                        let gen_sig = this.infcx.generator_sig(def_id).unwrap();
+                        // FIXME: want to revise this strategy when generator
+                        // signatures can actually contain LBRs.
+                        let sig = this.tcx().no_late_bound_regions(&gen_sig)
+                            .unwrap_or_else(|| {
+                                bug!("late-bound regions in signature of {:?}",
+                                     def_id)
+                            });
+                        (sig.yield_ty, sig.return_ty).fold_with(this)
+                    },
+                    |substs| {
+                        tcx.mk_generator(def_id, ty::ClosureSubsts { substs }, interior)
+                    }
+                )
+            }
+
             ty::TyBool |
             ty::TyChar |
             ty::TyInt(..) |
@@ -165,8 +309,6 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
             ty::TyFnDef(..) |
             ty::TyFnPtr(_) |
             ty::TyDynamic(..) |
-            ty::TyClosure(..) |
-            ty::TyGenerator(..) |
             ty::TyNever |
             ty::TyTuple(..) |
             ty::TyProjection(..) |
diff --git a/src/librustc/traits/select.rs b/src/librustc/traits/select.rs
index 44b8af3c1df..d6f27699e0a 100644
--- a/src/librustc/traits/select.rs
+++ b/src/librustc/traits/select.rs
@@ -904,14 +904,9 @@ impl<'cx, 'gcx, 'tcx> SelectionContext<'cx, 'gcx, 'tcx> {
                                dep_node: DepNodeIndex,
                                result: EvaluationResult)
     {
-        // Avoid caching results that depend on more than just the trait-ref:
-        // The stack can create recursion, and closure signatures
-        // being yet uninferred can create "spurious" EvaluatedToAmbig
-        // and EvaluatedToOk.
-        if result.is_stack_dependent() ||
-            ((result == EvaluatedToAmbig || result == EvaluatedToOk)
-             && trait_ref.has_closure_types())
-        {
+        // Avoid caching results that depend on more than just the trait-ref
+        // - the stack can create recursion.
+        if result.is_stack_dependent() {
             return;
         }
 
@@ -971,15 +966,12 @@ impl<'cx, 'gcx, 'tcx> SelectionContext<'cx, 'gcx, 'tcx> {
             this.candidate_from_obligation_no_cache(stack)
         });
 
-        if self.should_update_candidate_cache(&cache_fresh_trait_pred, &candidate) {
-            debug!("CACHE MISS: SELECT({:?})={:?}",
-                   cache_fresh_trait_pred, candidate);
-            self.insert_candidate_cache(stack.obligation.param_env,
-                                        cache_fresh_trait_pred,
-                                        dep_node,
-                                        candidate.clone());
-        }
-
+        debug!("CACHE MISS: SELECT({:?})={:?}",
+               cache_fresh_trait_pred, candidate);
+        self.insert_candidate_cache(stack.obligation.param_env,
+                                    cache_fresh_trait_pred,
+                                    dep_node,
+                                    candidate.clone());
         candidate
     }
 
@@ -1219,45 +1211,6 @@ impl<'cx, 'gcx, 'tcx> SelectionContext<'cx, 'gcx, 'tcx> {
                                   .insert(trait_ref, WithDepNode::new(dep_node, candidate));
     }
 
-    fn should_update_candidate_cache(&mut self,
-                                     cache_fresh_trait_pred: &ty::PolyTraitPredicate<'tcx>,
-                                     candidate: &SelectionResult<'tcx, SelectionCandidate<'tcx>>)
-                                     -> bool
-    {
-        // In general, it's a good idea to cache results, even
-        // ambiguous ones, to save us some trouble later. But we have
-        // to be careful not to cache results that could be
-        // invalidated later by advances in inference. Normally, this
-        // is not an issue, because any inference variables whose
-        // types are not yet bound are "freshened" in the cache key,
-        // which means that if we later get the same request once that
-        // type variable IS bound, we'll have a different cache key.
-        // For example, if we have `Vec<_#0t> : Foo`, and `_#0t` is
-        // not yet known, we may cache the result as `None`. But if
-        // later `_#0t` is bound to `Bar`, then when we freshen we'll
-        // have `Vec<Bar> : Foo` as the cache key.
-        //
-        // HOWEVER, it CAN happen that we get an ambiguity result in
-        // one particular case around closures where the cache key
-        // would not change. That is when the precise types of the
-        // upvars that a closure references have not yet been figured
-        // out (i.e., because it is not yet known if they are captured
-        // by ref, and if by ref, what kind of ref). In these cases,
-        // when matching a builtin bound, we will yield back an
-        // ambiguous result. But the *cache key* is just the closure type,
-        // it doesn't capture the state of the upvar computation.
-        //
-        // To avoid this trap, just don't cache ambiguous results if
-        // the self-type contains no inference byproducts (that really
-        // shouldn't happen in other circumstances anyway, given
-        // coherence).
-
-        match *candidate {
-            Ok(Some(_)) | Err(_) => true,
-            Ok(None) => cache_fresh_trait_pred.has_infer_types()
-        }
-    }
-
     fn assemble_candidates<'o>(&mut self,
                                stack: &TraitObligationStack<'o, 'tcx>)
                                -> Result<SelectionCandidateSet<'tcx>, SelectionError<'tcx>>