Speed up valid hole-fits by adding early abort and checks.

By adding an early abort flag in `TcSEnv`, we can fail fast in the presence of
insoluble constraints. This helps us avoid a lot of work in valid hole-fits, and
we geta massive speed-up by avoiding a lot of useless work solving constraints that
never come into play.

Additionally, we add a simple check for degenerate hole types, such as
when the type of the hole is an immutable type variable (as is the case
when the hole is completely unconstrained). Then the only valid fits are
the locals, so we can ignore the global candidates.

This fixes #16875

3 files changed, 108 insertions, 61 deletions

diff --git a/compiler/GHC/Tc/Errors/Hole.hs b/compiler/GHC/Tc/Errors/Hole.hs
index 4945b973e2..00e948bd10 100644
--- a/compiler/GHC/Tc/Errors/Hole.hs
+++ b/compiler/GHC/Tc/Errors/Hole.hs
@@ -66,7 +66,7 @@ import Data.List        ( partition, sort, sortOn, nubBy )
 import Data.Graph       ( graphFromEdges, topSort )
 
 
-import GHC.Tc.Solver    ( simplifyTopWanteds, runTcSDeriveds )
+import GHC.Tc.Solver    ( simplifyTopWanteds, runTcSDerivedsEarlyAbort )
 import GHC.Tc.Utils.Unify ( tcSubTypeSigma )
 
 import GHC.HsToCore.Docs ( extractDocs )
@@ -391,6 +391,26 @@ cause bewildering error messages. The solution here is simple: if a candidate
 would cause the type checker to error, it is not a valid hole fit, and thus it
 is discarded.
 
+Note [Speeding up valid hole-fits]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+To fix #16875 we noted that a lot of time was being spent on uneccessary work.
+
+When we'd call `tcCheckHoleFit hole hole_ty ty`, we would end up by generating
+a constraint to show that `hole_ty ~ ty`, including any constraints in `ty`. For
+example, if `hole_ty = Int` and `ty = Foldable t => (a -> Bool) -> t a -> Bool`,
+we'd have `(a_a1pa[sk:1] -> Bool) -> t_t2jk[sk:1] a_a1pa[sk:1] -> Bool ~# Int`
+from the coercion, as well as `Foldable t_t2jk[sk:1]`. By adding a flag to
+`TcSEnv` and adding a `runTcSDerivedsEarlyAbort`, we can fail as soon as we hit
+an insoluble constraint. Since we don't need the result in the case that it
+fails, a boolean `False` (i.e. "it didn't work" from `runTcSDerivedsEarlyAbort`)
+is sufficient.
+
+We also check whether the type of the hole is an immutable type variable (i.e.
+a skolem). In that case, the only possible fits are fits of exactly that type,
+which can only come from the locals. This speeds things up quite a bit when we
+don't know anything about the type of the hole. This also helps with degenerate
+fits like (`id (_ :: a)` and `head (_ :: [a])`) when looking for fits of type
+`a`, where `a` is a skolem.
 -}
 
 data HoleFitDispConfig = HFDC { showWrap :: Bool
@@ -574,7 +594,11 @@ findValidHoleFits tidy_env implics simples h@(Hole { hole_sort = ExprHole _
                       map IdHFCand lclBinds ++ map GreHFCand lcl
            globals = map GreHFCand gbl
            syntax = map NameHFCand builtIns
-           to_check = locals ++ syntax ++ globals
+           -- If the hole is a rigid type-variable, then we only check the
+           -- locals, since only they can match the type (in a meaningful way).
+           only_locals = any isImmutableTyVar $ getTyVar_maybe hole_ty
+           to_check = if only_locals then locals
+                      else locals ++ syntax ++ globals
      ; cands <- foldM (flip ($)) to_check candidatePlugins
      ; traceTc "numPlugins are:" $ ppr (length candidatePlugins)
      ; (searchDiscards, subs) <-
@@ -876,7 +900,6 @@ tcFilterHoleFits limit typed_hole ht@(hole_ty, _) candidates =
          ; traceTc "Did it fit?" $ ppr fits
          ; traceTc "wrap is: " $ ppr wrp
          ; traceTc "checkingFitOf }" empty
-         ; z_wrp_tys <- zonkTcTypes (unfoldWrapper wrp)
          -- We'd like to avoid refinement suggestions like `id _ _` or
          -- `head _ _`, and only suggest refinements where our all phantom
          -- variables got unified during the checking. This can be disabled
@@ -885,24 +908,26 @@ tcFilterHoleFits limit typed_hole ht@(hole_ty, _) candidates =
          -- variables, i.e. zonk them to read their final value to check for
          -- abstract refinements, and to report what the type of the simulated
          -- holes must be for this to be a match.
-         ; if fits
-           then if null ref_vars
-                then return (Just (z_wrp_tys, []))
-                else do { let -- To be concrete matches, matches have to
-                              -- be more than just an invented type variable.
-                              fvSet = fvVarSet fvs
-                              notAbstract :: TcType -> Bool
-                              notAbstract t = case getTyVar_maybe t of
-                                                Just tv -> tv `elemVarSet` fvSet
-                                                _ -> True
-                              allConcrete = all notAbstract z_wrp_tys
-                        ; z_vars  <- zonkTcTyVars ref_vars
-                        ; let z_mtvs = mapMaybe tcGetTyVar_maybe z_vars
-                        ; allFilled <- not <$> anyM isFlexiTyVar z_mtvs
-                        ; allowAbstract <- goptM Opt_AbstractRefHoleFits
-                        ; if allowAbstract || (allFilled && allConcrete )
-                          then return $ Just (z_wrp_tys, z_vars)
-                          else return Nothing }
+         ; if fits then do {
+              -- Zonking is expensive, so we only do it if required.
+              z_wrp_tys <- zonkTcTypes (unfoldWrapper wrp)
+            ; if null ref_vars
+              then return (Just (z_wrp_tys, []))
+              else do { let -- To be concrete matches, matches have to
+                            -- be more than just an invented type variable.
+                            fvSet = fvVarSet fvs
+                            notAbstract :: TcType -> Bool
+                            notAbstract t = case getTyVar_maybe t of
+                                              Just tv -> tv `elemVarSet` fvSet
+                                              _ -> True
+                            allConcrete = all notAbstract z_wrp_tys
+                      ; z_vars  <- zonkTcTyVars ref_vars
+                      ; let z_mtvs = mapMaybe tcGetTyVar_maybe z_vars
+                      ; allFilled <- not <$> anyM isFlexiTyVar z_mtvs
+                      ; allowAbstract <- goptM Opt_AbstractRefHoleFits
+                      ; if allowAbstract || (allFilled && allConcrete )
+                        then return $ Just (z_wrp_tys, z_vars)
+                        else return Nothing }}
            else return Nothing }
      where fvs = mkFVs ref_vars `unionFV` hole_fvs `unionFV` tyCoFVsOfType ty
            hole = typed_hole { th_hole = Nothing }
@@ -942,7 +967,8 @@ tcSubsumes ty_a ty_b = fst <$> tcCheckHoleFit dummyHole ty_a ty_b
 -- constraints on the type of the hole.
 tcCheckHoleFit :: TypedHole   -- ^ The hole to check against
                -> TcSigmaType
-               -- ^ The type to check against (possibly modified, e.g. refined)
+               -- ^ The type of the hole to check against (possibly modified,
+               -- e.g. refined with additional holes for refinement hole-fits.)
                -> TcSigmaType -- ^ The type to check whether fits.
                -> TcM (Bool, HsWrapper)
                -- ^ Whether it was a match, and the wrapper from hole_ty to ty.
@@ -970,22 +996,21 @@ tcCheckHoleFit (TypedHole {..}) hole_ty ty = discardErrs $
                 -- The relevant constraints may contain HoleDests, so we must
                 -- take care to clone them as well (to avoid #15370).
                ; cloned_relevants <- mapBagM cloneWanted th_relevant_cts
-                 -- We wrap the WC in the nested implications, see
+                 -- We wrap the WC in the nested implications, for details, see
                  -- Note [Checking hole fits]
-               ; let outermost_first = reverse th_implics
-                    -- We add the cloned relevants to the wanteds generated by
-                    -- the call to tcSubType_NC, see Note [Relevant constraints]
-                    -- There's no need to clone the wanteds, because they are
-                    -- freshly generated by `tcSubtype_NC`.
-                     w_rel_cts = addSimples wanted cloned_relevants
-                     final_wc  = foldr (setWCAndBinds fresh_binds) w_rel_cts outermost_first
+               ; let wrapInImpls cts = foldl (flip (setWCAndBinds fresh_binds)) cts th_implics
+                     final_wc  = wrapInImpls $ addSimples wanted cloned_relevants
+                 -- We add the cloned relevants to the wanteds generated
+                 -- by the call to tcSubType_NC, for details, see
+                 -- Note [Relevant constraints]. There's no need to clone
+                 -- the wanteds, because they are freshly generated by the
+                 -- call to`tcSubtype_NC`.
                ; traceTc "final_wc is: " $ ppr final_wc
-               ; rem <- runTcSDeriveds $ simplifyTopWanteds final_wc
-               -- We don't want any insoluble or simple constraints left, but
-               -- solved implications are ok (and necessary for e.g. undefined)
-               ; traceTc "rems was:" $ ppr rem
+                 -- See Note [Speeding up valid-hole fits]
+               ; (rem, _) <- tryTc $ runTcSDerivedsEarlyAbort $ simplifyTopWanteds final_wc
                ; traceTc "}" empty
-               ; return (isSolvedWC rem, wrap) } }
+               ; return (any isSolvedWC rem, wrap)
+               } }
      where
        setWCAndBinds :: EvBindsVar         -- Fresh ev binds var.
                      -> Implication        -- The implication to put WC in.
diff --git a/compiler/GHC/Tc/Solver.hs b/compiler/GHC/Tc/Solver.hs
index 6d9770adbf..276c0b284b 100644
--- a/compiler/GHC/Tc/Solver.hs
+++ b/compiler/GHC/Tc/Solver.hs
@@ -23,7 +23,11 @@ module GHC.Tc.Solver(
 
        -- For Rules we need these
        solveWanteds, solveWantedsAndDrop,
-       approximateWC, runTcSDeriveds
+       approximateWC, runTcSDeriveds,
+
+       -- We need this for valid hole-fits
+       runTcSDerivedsEarlyAbort
+
   ) where
 
 import GHC.Prelude
diff --git a/compiler/GHC/Tc/Solver/Monad.hs b/compiler/GHC/Tc/Solver/Monad.hs
index b957b0ed0c..7baf9ea186 100644
--- a/compiler/GHC/Tc/Solver/Monad.hs
+++ b/compiler/GHC/Tc/Solver/Monad.hs
@@ -13,9 +13,8 @@
 module GHC.Tc.Solver.Monad (
 
     -- The TcS monad
-    TcS, runTcS, runTcSDeriveds, runTcSWithEvBinds, runTcSInerts,
-    failTcS, warnTcS, addErrTcS, wrapTcS,
-    runTcSEqualities,
+    TcS, runTcS, runTcSDeriveds, runTcSDerivedsEarlyAbort, runTcSWithEvBinds,
+    runTcSInerts, failTcS, warnTcS, addErrTcS, wrapTcS, runTcSEqualities,
     nestTcS, nestImplicTcS, setEvBindsTcS,
     emitImplicationTcS, emitTvImplicationTcS,
 
@@ -614,10 +613,11 @@ When adding an equality to the inerts:
 addInertCan :: Ct -> TcS ()
 -- Precondition: item /is/ canonical
 -- See Note [Adding an equality to the InertCans]
-addInertCan ct
-  = do { traceTcS "addInertCan {" $
+addInertCan ct =
+    do { traceTcS "addInertCan {" $
          text "Trying to insert new inert item:" <+> ppr ct
-
+       ; mkTcS (\TcSEnv{tcs_abort_on_insoluble=abort_flag} ->
+                 when (abort_flag && insolubleEqCt ct) TcM.failM)
        ; ics <- getInertCans
        ; ct  <- maybeEmitShadow ics ct
        ; ics <- maybeKickOut ics ct
@@ -1198,6 +1198,11 @@ data TcSEnv
 
       tcs_inerts    :: IORef InertSet, -- Current inert set
 
+      -- Whether to throw an exception if we come across an insoluble constraint.
+      -- Used to fail-fast when checking for hole-fits. See Note [Speeding up
+      -- valid hole-fits].
+      tcs_abort_on_insoluble :: Bool,
+
       -- See Note [WorkList priorities] in GHC.Tc.Solver.InertSet
       tcs_worklist  :: IORef WorkList -- Current worklist
     }
@@ -1313,6 +1318,7 @@ runTcS tcs
        ; res <- runTcSWithEvBinds ev_binds_var tcs
        ; ev_binds <- TcM.getTcEvBindsMap ev_binds_var
        ; return (res, ev_binds) }
+
 -- | This variant of 'runTcS' will keep solving, even when only Deriveds
 -- are left around. It also doesn't return any evidence, as callers won't
 -- need it.
@@ -1321,6 +1327,14 @@ runTcSDeriveds tcs
   = do { ev_binds_var <- TcM.newTcEvBinds
        ; runTcSWithEvBinds ev_binds_var tcs }
 
+
+-- | This variant of 'runTcSDeriveds' will immediatley fail upon encountering an
+-- insoluble ct. See Note [Speeding up valid-hole fits]
+runTcSDerivedsEarlyAbort :: TcS a -> TcM a
+runTcSDerivedsEarlyAbort tcs
+  = do { ev_binds_var <- TcM.newTcEvBinds
+       ; runTcSWithEvBinds' True True ev_binds_var tcs }
+
 -- | This can deal only with equality constraints.
 runTcSEqualities :: TcS a -> TcM a
 runTcSEqualities thing_inside
@@ -1332,7 +1346,7 @@ runTcSEqualities thing_inside
 runTcSInerts :: InertSet -> TcS a -> TcM (a, InertSet)
 runTcSInerts inerts tcs = do
   ev_binds_var <- TcM.newTcEvBinds
-  runTcSWithEvBinds' False ev_binds_var $ do
+  runTcSWithEvBinds' False False ev_binds_var $ do
     setTcSInerts inerts
     a <- tcs
     new_inerts <- getTcSInerts
@@ -1341,27 +1355,29 @@ runTcSInerts inerts tcs = do
 runTcSWithEvBinds :: EvBindsVar
                   -> TcS a
                   -> TcM a
-runTcSWithEvBinds = runTcSWithEvBinds' True
+runTcSWithEvBinds = runTcSWithEvBinds' True False
 
 runTcSWithEvBinds' :: Bool -- ^ Restore type variable cycles afterwards?
                            -- Don't if you want to reuse the InertSet.
                            -- See also Note [Type variable cycles]
                            -- in GHC.Tc.Solver.Canonical
+                   -> Bool
                    -> EvBindsVar
                    -> TcS a
                    -> TcM a
-runTcSWithEvBinds' restore_cycles ev_binds_var tcs
+runTcSWithEvBinds' restore_cycles abort_on_insoluble ev_binds_var tcs
   = do { unified_var <- TcM.newTcRef 0
        ; step_count <- TcM.newTcRef 0
        ; inert_var <- TcM.newTcRef emptyInert
        ; wl_var <- TcM.newTcRef emptyWorkList
        ; unif_lvl_var <- TcM.newTcRef Nothing
-       ; let env = TcSEnv { tcs_ev_binds      = ev_binds_var
-                          , tcs_unified       = unified_var
-                          , tcs_unif_lvl      = unif_lvl_var
-                          , tcs_count         = step_count
-                          , tcs_inerts        = inert_var
-                          , tcs_worklist      = wl_var }
+       ; let env = TcSEnv { tcs_ev_binds           = ev_binds_var
+                          , tcs_unified            = unified_var
+                          , tcs_unif_lvl           = unif_lvl_var
+                          , tcs_count              = step_count
+                          , tcs_inerts             = inert_var
+                          , tcs_abort_on_insoluble = abort_on_insoluble
+                          , tcs_worklist           = wl_var }
 
              -- Run the computation
        ; res <- unTcS tcs env
@@ -1418,10 +1434,11 @@ nestImplicTcS :: EvBindsVar
               -> TcLevel -> TcS a
               -> TcS a
 nestImplicTcS ref inner_tclvl (TcS thing_inside)
-  = TcS $ \ TcSEnv { tcs_unified       = unified_var
-                   , tcs_inerts        = old_inert_var
-                   , tcs_count         = count
-                   , tcs_unif_lvl      = unif_lvl
+  = TcS $ \ TcSEnv { tcs_unified            = unified_var
+                   , tcs_inerts             = old_inert_var
+                   , tcs_count              = count
+                   , tcs_unif_lvl           = unif_lvl
+                   , tcs_abort_on_insoluble = abort_on_insoluble
                    } ->
     do { inerts <- TcM.readTcRef old_inert_var
        ; let nest_inert = inerts { inert_cycle_breakers = []
@@ -1430,12 +1447,13 @@ nestImplicTcS ref inner_tclvl (TcS thing_inside)
                  -- All other InertSet fields are inherited
        ; new_inert_var <- TcM.newTcRef nest_inert
        ; new_wl_var    <- TcM.newTcRef emptyWorkList
-       ; let nest_env = TcSEnv { tcs_count         = count     -- Inherited
-                               , tcs_unif_lvl      = unif_lvl  -- Inherited
-                               , tcs_ev_binds      = ref
-                               , tcs_unified       = unified_var
-                               , tcs_inerts        = new_inert_var
-                               , tcs_worklist      = new_wl_var }
+       ; let nest_env = TcSEnv { tcs_count              = count     -- Inherited
+                               , tcs_unif_lvl           = unif_lvl  -- Inherited
+                               , tcs_ev_binds           = ref
+                               , tcs_unified            = unified_var
+                               , tcs_inerts             = new_inert_var
+                               , tcs_abort_on_insoluble = abort_on_insoluble
+                               , tcs_worklist           = new_wl_var }
        ; res <- TcM.setTcLevel inner_tclvl $
                 thing_inside nest_env