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 rust-lang#43787.

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

Author: arielb1

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
@@ -33,6 +44,8 @@
 use ty::{self, Ty, TyCtxt, TypeFoldable};
 use ty::fold::TypeFolder;
 use util::nodemap::FxHashMap;
+use hir::def_id::DefId;
+
 use std::collections::hash_map::Entry;
 
 use super::InferCtxt;
@@ -42,6 +55,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 +65,7 @@ impl<'a, 'gcx, 'tcx> TypeFreshener<'a, 'gcx, 'tcx> {
             infcx,
             freshen_count: 0,
             freshen_map: FxHashMap(),
+            closure_set: vec![],
         }
     }
 
@@ -76,6 +91,16 @@ 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))
+    }
 }
 
 impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
@@ -105,7 +130,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 +176,82 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
                 t
             }
 
+            ty::TyClosure(def_id, substs) => {
+                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 closure_sig_marker = 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)
+                } else {
+                    self.closure_set.push(def_id);
+                    let closure_sig = self.infcx.fn_sig(def_id);
+                    let closure_sig_marker = tcx.mk_fn_ptr(closure_sig.fold_with(self));
+                    self.closure_set.pop();
+                    closure_sig_marker
+                };
+
+                // 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 = self.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 params = tcx.mk_substs(
+                    substs.substs.iter().map(|k| k.fold_with(self)).chain(
+                        [closure_sig_marker, closure_kind_marker].iter().cloned().map(From::from)
+                    ));
+                tcx.mk_closure(def_id, params)
+            }
+
             ty::TyBool |
             ty::TyChar |
             ty::TyInt(..) |
@@ -165,7 +267,6 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for TypeFreshener<'a, 'gcx, 'tcx> {
             ty::TyFnDef(..) |
             ty::TyFnPtr(_) |
             ty::TyDynamic(..) |
-            ty::TyClosure(..) |
             ty::TyNever |
             ty::TyTuple(..) |
             ty::TyProjection(..) |

src/librustc/traits/select.rs

@@ -888,14 +888,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;
         }
 
@@ -955,15 +950,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
     }
 
@@ -1203,45 +1195,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>>