When flattening, try reducing type-family applications eagerly

This short-cut can improve performance quite a bit, by short-circuiting
the process of creating a fresh constraint and binding for each reduction.

See Note [Reduce type family applications eagerly] in TcFlatten

To do this I had to generalise the inert_flat_cache a bit, so that the
rhs is not necessarily a type variable; but nothing fundamental.

3 files changed, 50 insertions, 18 deletions

diff --git a/compiler/typecheck/TcCanonical.hs b/compiler/typecheck/TcCanonical.hs
index 669dc06672..4042fe83c9 100644
--- a/compiler/typecheck/TcCanonical.hs
+++ b/compiler/typecheck/TcCanonical.hs
@@ -649,7 +649,7 @@ canCFunEqCan ev fn tys fsk
            Stop ev s        -> return (Stop ev s) ;
            ContinueWith ev' ->
 
-    do { extendFlatCache fn tys' (ctEvCoercion ev', fsk)
+    do { extendFlatCache fn tys' (ctEvCoercion ev', fsk_ty, ev')
        ; continueWith (CFunEqCan { cc_ev = ev', cc_fun = fn
                                  , cc_tyargs = tys', cc_fsk = fsk }) } } }
 
diff --git a/compiler/typecheck/TcFlatten.hs b/compiler/typecheck/TcFlatten.hs
index 6ab8b227d0..f8d2148dd7 100644
--- a/compiler/typecheck/TcFlatten.hs
+++ b/compiler/typecheck/TcFlatten.hs
@@ -798,20 +798,32 @@ flatten_exact_fam_app_fully fmode tc tys
 
        ; mb_ct <- lookupFlatCache tc xis
        ; case mb_ct of
-           Just (co, fsk)  -- co :: F xis ~ fsk
-             | isFskTyVar fsk || not (isGiven ctxt_ev)
+           Just (co, rhs_ty, ev)  -- co :: F xis ~ fsk
+             | ev `canRewriteOrSame` ctxt_ev
              ->  -- Usable hit in the flat-cache
-                 -- isFskTyVar checks for a "given" in the cache
-                do { traceTcS "flatten/flat-cache hit" $ (ppr tc <+> ppr xis $$ ppr fsk $$ ppr co)
-                   ; (fsk_xi, fsk_co) <- flattenTyVar fmode fsk
+                 -- We certainly *can* use a Wanted for a Wanted
+                do { traceTcS "flatten/flat-cache hit" $ (ppr tc <+> ppr xis $$ ppr rhs_ty $$ ppr co)
+                   ; (fsk_xi, fsk_co) <- flatten_one fmode rhs_ty
                           -- The fsk may already have been unified, so flatten it
                           -- fsk_co :: fsk_xi ~ fsk
                    ; return (fsk_xi, fsk_co `mkTcTransCo` mkTcSymCo co `mkTcTransCo` ret_co) }
                                     -- :: fsk_xi ~ F xis
 
-           _ -> do { let fam_ty = mkTyConApp tc xis
+           -- Try to reduce the family application right now
+           -- See Note [Reduce type family applications eagerly]
+           _ -> do { mb_match <- matchFam tc xis
+                   ; case mb_match of {
+                        Just (norm_co, norm_ty)
+                            -> do { (xi, final_co) <- flatten_one fmode norm_ty
+                                  ; let co = norm_co `mkTcTransCo` mkTcSymCo final_co
+                                  ; extendFlatCache tc xis (co, xi, ctxt_ev)
+                                  ; return (xi, mkTcSymCo co `mkTcTransCo` ret_co) } ;
+                        Nothing ->
+                do { let fam_ty = mkTyConApp tc xis
                    ; (ev, fsk) <- newFlattenSkolem ctxt_ev fam_ty
-                   ; extendFlatCache tc xis (ctEvCoercion ev, fsk)
+                   ; let fsk_ty = mkTyVarTy fsk
+                         co     = ctEvCoercion ev
+                   ; extendFlatCache tc xis (co, fsk_ty, ev)
 
                    -- The new constraint (F xis ~ fsk) is not necessarily inert
                    -- (e.g. the LHS may be a redex) so we must put it in the work list
@@ -822,9 +834,26 @@ flatten_exact_fam_app_fully fmode tc tys
                    ; emitFlatWork ct
 
                    ; traceTcS "flatten/flat-cache miss" $ (ppr fam_ty $$ ppr fsk $$ ppr ev)
-                   ; return (mkTyVarTy fsk, mkTcSymCo (ctEvCoercion ev) `mkTcTransCo` ret_co) } }
+                   ; return (fsk_ty, mkTcSymCo co `mkTcTransCo` ret_co) }
+        } } }
+
+{- Note [Reduce type family applications eagerly]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+If we come across a type-family application like (Append (Cons x Nil) t),
+then, rather than flattening to a skolem etc, we may as well just reduce
+it on the spot to (Cons x t).  This saves a lot of intermediate steps.
+Examples that are helped are tests T9872, and T5321Fun.
+
+So just before we create the new skolem, we attempt to reduce it by one
+step (using matchFam).  If that works, then recursively flatten the rhs,
+which may in turn do lots more reductions.
+
+Once we've got a flat rhs, we extend the flatten-cache to record the
+result.  Doing so can save lots of work when the same redex shows up
+more than once.  Note that we record the link from the redex all the
+way to its *final* value, not just the single step reduction.
+
 
-{-
 ************************************************************************
 *                                                                      *
              Flattening a type variable
diff --git a/compiler/typecheck/TcSMonad.hs b/compiler/typecheck/TcSMonad.hs
index 204a471f9c..cba8e24dd6 100644
--- a/compiler/typecheck/TcSMonad.hs
+++ b/compiler/typecheck/TcSMonad.hs
@@ -430,10 +430,13 @@ data InertSet
               -- Canonical Given, Wanted, Derived (no Solved)
               -- Sometimes called "the inert set"
 
-       , inert_flat_cache :: FunEqMap (TcCoercion, TcTyVar)
+       , inert_flat_cache :: FunEqMap (TcCoercion, TcType, CtEvidence)
               -- See Note [Type family equations]
-              -- If    F tys :-> (co, fsk),
-              -- then  co :: F tys ~ fsk
+              -- If    F tys :-> (co, ty, ev),
+              -- then  co :: F tys ~ ty
+              --
+              -- The 'ev' field is just for the G/W/D flavour, nothing more!
+              --
               -- Just a hash-cons cache for use when flattening only
               -- These include entirely un-processed goals, so don't use
               -- them to solve a top-level goal, else you may end up solving
@@ -799,7 +802,7 @@ checkAllSolved
                      || unsolved_dicts || unsolved_funeqs
                      || not (isEmptyBag (inert_insols icans)))) }
 
-lookupFlatCache :: TyCon -> [Type] -> TcS (Maybe (TcCoercion, TcTyVar))
+lookupFlatCache :: TyCon -> [Type] -> TcS (Maybe (TcCoercion, TcType, CtEvidence))
 lookupFlatCache fam_tc tys
   = do { IS { inert_flat_cache = flat_cache
             , inert_cans = IC { inert_funeqs = inert_funeqs } } <- getTcSInerts
@@ -809,7 +812,7 @@ lookupFlatCache fam_tc tys
     lookup_inerts inert_funeqs
       | Just (CFunEqCan { cc_ev = ctev, cc_fsk = fsk })
            <- findFunEqs inert_funeqs fam_tc tys
-      = Just (ctEvCoercion ctev, fsk)
+      = Just (ctEvCoercion ctev, mkTyVarTy fsk, ctev)
       | otherwise = Nothing
 
     lookup_flats flat_cache = findFunEq flat_cache fam_tc tys
@@ -1546,12 +1549,12 @@ newFlattenSkolem ctxt_ev fam_ty
   where
     loc = ctEvLoc ctxt_ev
 
-extendFlatCache :: TyCon -> [Type] -> (TcCoercion, TcTyVar) -> TcS ()
-extendFlatCache tc xi_args (co, fsk)
+extendFlatCache :: TyCon -> [Type] -> (TcCoercion, TcType, CtEvidence) -> TcS ()
+extendFlatCache tc xi_args stuff
   = do { dflags <- getDynFlags
        ; when (gopt Opt_FlatCache dflags) $
          updInertTcS $ \ is@(IS { inert_flat_cache = fc }) ->
-            is { inert_flat_cache = insertFunEq fc tc xi_args (co, fsk) } }
+            is { inert_flat_cache = insertFunEq fc tc xi_args stuff } }
 
 -- Instantiations
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~