compiler: Introduce and use RoughMap for instance environmentswip/roughmap-mp

Here we introduce a new data structure, RoughMap, inspired by the
previous `RoughTc` matching mechanism for checking instance matches.
This allows [Fam]InstEnv to be implemented as a trie indexed by these
RoughTc signatures, reducing the complexity of instance lookup and
FamInstEnv merging (done during the family instance conflict test)
from O(n) to O(log n).

The critical performance improvement currently realised by this patch is
in instance matching. In particular the RoughMap mechanism allows us to
discount many potential instances which will never match for constraints
involving type variables (see Note [Matching a RoughMap]). In realistic
code bases matchInstEnv was accounting for 50% of typechecker time due
to redundant work checking instances when simplifying instance contexts
when deriving instances. With this patch the cost is significantly
reduced.

The larger constants in InstEnv creation do mean that a few small
tests regress in allocations slightly. However, the runtime of T19703 is
reduced by a factor of 4. Moreover, the compilation time of the Cabal
library is slightly improved.

A couple of test cases are included which demonstrate significant
improvements in compile time with this patch.

This unfortunately does not fix the testcase provided in #19703 but does
fix #20933

-------------------------
Metric Decrease:
    T12425
Metric Increase:
    T13719
    T9872a
    T9872d
    hard_hole_fits
-------------------------

Co-authored-by: Matthew Pickering <matthewtpickering@gmail.com>

1 files changed, 125 insertions, 78 deletions

diff --git a/compiler/GHC/Core/InstEnv.hs b/compiler/GHC/Core/InstEnv.hs
index ab23fcae2c..e223a7cd87 100644
--- a/compiler/GHC/Core/InstEnv.hs
+++ b/compiler/GHC/Core/InstEnv.hs
@@ -11,17 +11,19 @@ The bits common to GHC.Tc.TyCl.Instance and GHC.Tc.Deriv.
 
 module GHC.Core.InstEnv (
         DFunId, InstMatch, ClsInstLookupResult,
+        PotentialUnifiers(..), getPotentialUnifiers, nullUnifiers,
         OverlapFlag(..), OverlapMode(..), setOverlapModeMaybe,
         ClsInst(..), DFunInstType, pprInstance, pprInstanceHdr, pprInstances,
         instanceHead, instanceSig, mkLocalInstance, mkImportedInstance,
-        instanceDFunId, updateClsInstDFun, instanceRoughTcs,
+        instanceDFunId, updateClsInstDFuns, updateClsInstDFun,
         fuzzyClsInstCmp, orphNamesOfClsInst,
 
         InstEnvs(..), VisibleOrphanModules, InstEnv,
-        emptyInstEnv, extendInstEnv,
-        deleteFromInstEnv, deleteDFunFromInstEnv,
+        mkInstEnv, emptyInstEnv, unionInstEnv, extendInstEnv,
+        filterInstEnv, deleteFromInstEnv, deleteDFunFromInstEnv,
+        anyInstEnv,
         identicalClsInstHead,
-        extendInstEnvList, lookupUniqueInstEnv, lookupInstEnv, instEnvElts, instEnvClasses,
+        extendInstEnvList, lookupUniqueInstEnv, lookupInstEnv, instEnvElts, instEnvClasses, mapInstEnv,
         memberInstEnv,
         instIsVisible,
         classInstances, instanceBindFun,
@@ -34,25 +36,25 @@ import GHC.Prelude
 import GHC.Tc.Utils.TcType -- InstEnv is really part of the type checker,
               -- and depends on TcType in many ways
 import GHC.Core ( IsOrphan(..), isOrphan, chooseOrphanAnchor )
+import GHC.Core.RoughMap
 import GHC.Unit.Module.Env
 import GHC.Unit.Types
 import GHC.Core.Class
 import GHC.Types.Var
+import GHC.Types.Unique.DSet
 import GHC.Types.Var.Set
 import GHC.Types.Name
 import GHC.Types.Name.Set
-import GHC.Types.Unique (getUnique)
 import GHC.Core.Unify
 import GHC.Types.Basic
-import GHC.Types.Unique.DFM
 import GHC.Types.Id
 import Data.Data        ( Data )
 import Data.Maybe       ( isJust )
 
-import GHC.Utils.Misc
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Utils.Panic.Plain
+import Data.Semigroup
 
 {-
 ************************************************************************
@@ -68,9 +70,12 @@ import GHC.Utils.Panic.Plain
 data ClsInst
   = ClsInst {   -- Used for "rough matching"; see
                 -- Note [ClsInst laziness and the rough-match fields]
-                -- INVARIANT: is_tcs = roughMatchTcs is_tys
+                -- INVARIANT: is_tcs = KnownTc is_cls_nm : roughMatchTcs is_tys
                is_cls_nm :: Name          -- ^ Class name
+
              , is_tcs  :: [RoughMatchTc]  -- ^ Top of type args
+                          -- The class itself is always
+                          -- the first element of this list
 
                -- | @is_dfun_name = idName . is_dfun@.
                --
@@ -103,13 +108,12 @@ data ClsInst
 -- instances before displaying them to the user.
 fuzzyClsInstCmp :: ClsInst -> ClsInst -> Ordering
 fuzzyClsInstCmp x y =
-    stableNameCmp (is_cls_nm x) (is_cls_nm y) `mappend`
-    mconcat (map cmp (zip (is_tcs x) (is_tcs y)))
+    foldMap cmp (zip (is_tcs x) (is_tcs y))
   where
-    cmp (OtherTc,   OtherTc)   = EQ
-    cmp (OtherTc,   KnownTc _) = LT
-    cmp (KnownTc _, OtherTc)   = GT
-    cmp (KnownTc x, KnownTc y) = stableNameCmp x y
+    cmp (RM_WildCard,  RM_WildCard)   = EQ
+    cmp (RM_WildCard,  RM_KnownTc _) = LT
+    cmp (RM_KnownTc _, RM_WildCard)   = GT
+    cmp (RM_KnownTc x, RM_KnownTc y) = stableNameCmp x y
 
 isOverlappable, isOverlapping, isIncoherent :: ClsInst -> Bool
 isOverlappable i = hasOverlappableFlag (overlapMode (is_flag i))
@@ -196,8 +200,9 @@ updateClsInstDFun :: (DFunId -> DFunId) -> ClsInst -> ClsInst
 updateClsInstDFun tidy_dfun ispec
   = ispec { is_dfun = tidy_dfun (is_dfun ispec) }
 
-instanceRoughTcs :: ClsInst -> [RoughMatchTc]
-instanceRoughTcs = is_tcs
+updateClsInstDFuns :: (DFunId -> DFunId) -> InstEnv -> InstEnv
+updateClsInstDFuns tidy_dfun (InstEnv rm)
+  = InstEnv $ fmap (updateClsInstDFun tidy_dfun) rm
 
 instance NamedThing ClsInst where
    getName ispec = getName (is_dfun ispec)
@@ -259,7 +264,7 @@ mkLocalInstance dfun oflag tvs cls tys
             , is_tvs = tvs
             , is_dfun_name = dfun_name
             , is_cls = cls, is_cls_nm = cls_name
-            , is_tys = tys, is_tcs = roughMatchTcs tys
+            , is_tys = tys, is_tcs = RM_KnownTc cls_name : roughMatchTcs tys
             , is_orphan = orph
             }
   where
@@ -290,7 +295,7 @@ mkLocalInstance dfun oflag tvs cls tys
     choose_one nss = chooseOrphanAnchor (unionNameSets nss)
 
 mkImportedInstance :: Name           -- ^ the name of the class
-                   -> [RoughMatchTc] -- ^ the types which the class was applied to
+                   -> [RoughMatchTc] -- ^ the rough match signature of the instance
                    -> Name           -- ^ the 'Name' of the dictionary binding
                    -> DFunId         -- ^ the 'Id' of the dictionary.
                    -> OverlapFlag    -- ^ may this instance overlap?
@@ -304,7 +309,8 @@ mkImportedInstance cls_nm mb_tcs dfun_name dfun oflag orphan
   = ClsInst { is_flag = oflag, is_dfun = dfun
             , is_tvs = tvs, is_tys = tys
             , is_dfun_name = dfun_name
-            , is_cls_nm = cls_nm, is_cls = cls, is_tcs = mb_tcs
+            , is_cls_nm = cls_nm, is_cls = cls
+            , is_tcs = RM_KnownTc cls_nm : mb_tcs
             , is_orphan = orphan }
   where
     (tvs, _, cls, tys) = tcSplitDFunTy (idType dfun)
@@ -386,9 +392,12 @@ UniqDFM. See also Note [Deterministic UniqFM]
 -- We still use Class as key type as it's both the common case
 -- and conveys the meaning better. But the implementation of
 --InstEnv is a bit more lax internally.
-type InstEnv = UniqDFM Class ClsInstEnv      -- Maps Class to instances for that class
+newtype InstEnv = InstEnv (RoughMap ClsInst)      -- Maps Class to instances for that class
   -- See Note [InstEnv determinism]
 
+instance Outputable InstEnv where
+  ppr (InstEnv rm) = pprInstances $ elemsRM rm
+
 -- | 'InstEnvs' represents the combination of the global type class instance
 -- environment, the local type class instance environment, and the set of
 -- transitively reachable orphan modules (according to what modules have been
@@ -406,30 +415,32 @@ data InstEnvs = InstEnvs {
 -- transitively reachable orphan modules (modules that define orphan instances).
 type VisibleOrphanModules = ModuleSet
 
-newtype ClsInstEnv
-  = ClsIE [ClsInst]    -- The instances for a particular class, in any order
-
-instance Outputable ClsInstEnv where
-  ppr (ClsIE is) = pprInstances is
 
 -- INVARIANTS:
 --  * The is_tvs are distinct in each ClsInst
 --      of a ClsInstEnv (so we can safely unify them)
 
--- Thus, the @ClassInstEnv@ for @Eq@ might contain the following entry:
+-- Thus, the @ClsInstEnv@ for @Eq@ might contain the following entry:
 --      [a] ===> dfun_Eq_List :: forall a. Eq a => Eq [a]
 -- The "a" in the pattern must be one of the forall'd variables in
 -- the dfun type.
 
 emptyInstEnv :: InstEnv
-emptyInstEnv = emptyUDFM
+emptyInstEnv = InstEnv emptyRM
+
+mkInstEnv :: [ClsInst] -> InstEnv
+mkInstEnv = extendInstEnvList emptyInstEnv
 
 instEnvElts :: InstEnv -> [ClsInst]
-instEnvElts ie = [elt | ClsIE elts <- eltsUDFM ie, elt <- elts]
+instEnvElts (InstEnv rm) = elemsRM rm
   -- See Note [InstEnv determinism]
 
-instEnvClasses :: InstEnv -> [Class]
-instEnvClasses ie = [is_cls e | ClsIE (e : _) <- eltsUDFM ie]
+instEnvEltsForClass :: InstEnv -> Class -> [ClsInst]
+instEnvEltsForClass (InstEnv rm) cls = lookupRM [RML_KnownTc (className cls)] rm
+
+-- N.B. this is not particularly efficient but used only by GHCi.
+instEnvClasses :: InstEnv -> UniqDSet Class
+instEnvClasses ie = mkUniqDSet $ map is_cls (instEnvElts ie)
 
 -- | Test if an instance is visible, by checking that its origin module
 -- is in 'VisibleOrphanModules'.
@@ -449,42 +460,50 @@ classInstances :: InstEnvs -> Class -> [ClsInst]
 classInstances (InstEnvs { ie_global = pkg_ie, ie_local = home_ie, ie_visible = vis_mods }) cls
   = get home_ie ++ get pkg_ie
   where
-    get env = case lookupUDFM env cls of
-                Just (ClsIE insts) -> filter (instIsVisible vis_mods) insts
-                Nothing            -> []
+    get :: InstEnv -> [ClsInst]
+    get ie = filter (instIsVisible vis_mods) (instEnvEltsForClass ie cls)
 
 -- | Checks for an exact match of ClsInst in the instance environment.
 -- We use this when we do signature checking in "GHC.Tc.Module"
 memberInstEnv :: InstEnv -> ClsInst -> Bool
-memberInstEnv inst_env ins_item@(ClsInst { is_cls_nm = cls_nm } ) =
-    maybe False (\(ClsIE items) -> any (identicalDFunType ins_item) items)
-          (lookupUDFM_Directly inst_env (getUnique cls_nm))
+memberInstEnv (InstEnv rm) ins_item@(ClsInst { is_tcs = tcs } ) =
+    any (identicalDFunType ins_item) (fst $ lookupRM' (map roughMatchTcToLookup tcs) rm)
  where
   identicalDFunType cls1 cls2 =
     eqType (varType (is_dfun cls1)) (varType (is_dfun cls2))
 
+-- | Makes no particular effort to detect conflicts.
+unionInstEnv :: InstEnv -> InstEnv -> InstEnv
+unionInstEnv (InstEnv a) (InstEnv b) = InstEnv (a `unionRM` b)
+
 extendInstEnvList :: InstEnv -> [ClsInst] -> InstEnv
 extendInstEnvList inst_env ispecs = foldl' extendInstEnv inst_env ispecs
 
 extendInstEnv :: InstEnv -> ClsInst -> InstEnv
-extendInstEnv inst_env ins_item@(ClsInst { is_cls_nm = cls_nm })
-  = addToUDFM_C_Directly add inst_env (getUnique cls_nm) (ClsIE [ins_item])
-  where
-    add (ClsIE cur_insts) _ = ClsIE (ins_item : cur_insts)
+extendInstEnv (InstEnv rm) ins_item@(ClsInst { is_tcs = tcs })
+  = InstEnv $ insertRM tcs ins_item rm
+
+filterInstEnv :: (ClsInst -> Bool) -> InstEnv -> InstEnv
+filterInstEnv pred (InstEnv rm)
+  = InstEnv $ filterRM pred rm
+
+anyInstEnv :: (ClsInst -> Bool) -> InstEnv -> Bool
+anyInstEnv pred (InstEnv rm)
+  = foldRM (\x rest -> pred x || rest) False rm
+
+mapInstEnv :: (ClsInst -> ClsInst) -> InstEnv -> InstEnv
+mapInstEnv f (InstEnv rm) = InstEnv (f <$> rm)
 
 deleteFromInstEnv :: InstEnv -> ClsInst -> InstEnv
-deleteFromInstEnv inst_env ins_item@(ClsInst { is_cls_nm = cls_nm })
-  = adjustUDFM_Directly adjust inst_env (getUnique cls_nm)
-  where
-    adjust (ClsIE items) = ClsIE (filterOut (identicalClsInstHead ins_item) items)
+deleteFromInstEnv (InstEnv rm) ins_item@(ClsInst { is_tcs = tcs })
+  = InstEnv $ filterMatchingRM (not . identicalClsInstHead ins_item) tcs rm
 
 deleteDFunFromInstEnv :: InstEnv -> DFunId -> InstEnv
 -- Delete a specific instance fron an InstEnv
-deleteDFunFromInstEnv inst_env dfun
-  = adjustUDFM adjust inst_env cls
+deleteDFunFromInstEnv (InstEnv rm) dfun
+  = InstEnv $ filterMatchingRM (not . same_dfun) [RM_KnownTc (className cls)] rm
   where
     (_, _, cls, _) = tcSplitDFunTy (idType dfun)
-    adjust (ClsIE items) = ClsIE (filterOut same_dfun items)
     same_dfun (ClsInst { is_dfun = dfun' }) = dfun == dfun'
 
 identicalClsInstHead :: ClsInst -> ClsInst -> Bool
@@ -492,10 +511,10 @@ identicalClsInstHead :: ClsInst -> ClsInst -> Bool
 -- e.g.  both are   Eq [(a,b)]
 -- Used for overriding in GHCi
 -- Obviously should be insensitive to alpha-renaming
-identicalClsInstHead (ClsInst { is_cls_nm = cls_nm1, is_tcs = rough1, is_tys = tys1 })
-                     (ClsInst { is_cls_nm = cls_nm2, is_tcs = rough2, is_tys = tys2 })
-  =  cls_nm1 == cls_nm2
-  && not (instanceCantMatch rough1 rough2)  -- Fast check for no match, uses the "rough match" fields
+identicalClsInstHead (ClsInst { is_tcs = rough1, is_tys = tys1 })
+                     (ClsInst { is_tcs = rough2, is_tys = tys2 })
+  =  not (instanceCantMatch rough1 rough2)  -- Fast check for no match, uses the "rough match" fields;
+                                            -- also accounts for class name.
   && isJust (tcMatchTys tys1 tys2)
   && isJust (tcMatchTys tys2 tys1)
 
@@ -730,7 +749,7 @@ type InstMatch = (ClsInst, [DFunInstType])
 
 type ClsInstLookupResult
      = ( [InstMatch]     -- Successful matches
-       , [ClsInst]       -- These don't match but do unify
+       , PotentialUnifiers  -- These don't match but do unify
        , [InstMatch] )   -- Unsafe overlapped instances under Safe Haskell
                          -- (see Note [Safe Haskell Overlapping Instances] in
                          -- GHC.Tc.Solver).
@@ -811,11 +830,38 @@ lookupUniqueInstEnv instEnv cls tys
       _other -> Left $ text "instance not found" <+>
                        (ppr $ mkTyConApp (classTyCon cls) tys)
 
+data PotentialUnifiers = NoUnifiers
+                       | OneOrMoreUnifiers [ClsInst]
+                       -- This list is lazy as we only look at all the unifiers when
+                       -- printing an error message. It can be expensive to compute all
+                       -- the unifiers because if you are matching something like C a[sk] then
+                       -- all instances will unify.
+
+instance Outputable PotentialUnifiers where
+  ppr NoUnifiers = text "NoUnifiers"
+  ppr xs = ppr (getPotentialUnifiers xs)
+
+instance Semigroup PotentialUnifiers where
+  NoUnifiers <> u = u
+  u <> NoUnifiers = u
+  u1 <> u2 = OneOrMoreUnifiers (getPotentialUnifiers u1 ++ getPotentialUnifiers u2)
+
+instance Monoid PotentialUnifiers where
+  mempty = NoUnifiers
+
+getPotentialUnifiers :: PotentialUnifiers -> [ClsInst]
+getPotentialUnifiers NoUnifiers = []
+getPotentialUnifiers (OneOrMoreUnifiers cls) = cls
+
+nullUnifiers :: PotentialUnifiers -> Bool
+nullUnifiers NoUnifiers = True
+nullUnifiers _ = False
+
 lookupInstEnv' :: InstEnv          -- InstEnv to look in
                -> VisibleOrphanModules   -- But filter against this
                -> Class -> [Type]  -- What we are looking for
                -> ([InstMatch],    -- Successful matches
-                   [ClsInst])      -- These don't match but do unify
+                   PotentialUnifiers)      -- These don't match but do unify
                                    -- (no incoherent ones in here)
 -- The second component of the result pair happens when we look up
 --      Foo [a]
@@ -827,35 +873,35 @@ lookupInstEnv' :: InstEnv          -- InstEnv to look in
 -- but Foo [Int] is a unifier.  This gives the caller a better chance of
 -- giving a suitable error message
 
-lookupInstEnv' ie vis_mods cls tys
-  = lookup ie
+lookupInstEnv' (InstEnv rm) vis_mods cls tys
+  = (foldr check_match [] rough_matches, check_unifier rough_unifiers)
   where
-    rough_tcs  = roughMatchTcs tys
-
-    --------------
-    lookup env = case lookupUDFM env cls of
-                   Nothing -> ([],[])   -- No instances for this class
-                   Just (ClsIE insts) -> find [] [] insts
+    (rough_matches, rough_unifiers) = lookupRM' rough_tcs rm
+    rough_tcs  = RML_KnownTc (className cls) : roughMatchTcsLookup tys
 
     --------------
-    find ms us [] = (ms, us)
-    find ms us (item@(ClsInst { is_tcs = mb_tcs, is_tvs = tpl_tvs
-                              , is_tys = tpl_tys }) : rest)
+    check_match :: ClsInst -> [InstMatch] -> [InstMatch]
+    check_match item@(ClsInst { is_tvs = tpl_tvs, is_tys = tpl_tys }) acc
       | not (instIsVisible vis_mods item)
-      = find ms us rest  -- See Note [Instance lookup and orphan instances]
-
-        -- Fast check for no match, uses the "rough match" fields
-      | instanceCantMatch rough_tcs mb_tcs
-      = find ms us rest
+      = acc  -- See Note [Instance lookup and orphan instances]
 
       | Just subst <- tcMatchTys tpl_tys tys
-      = find ((item, map (lookupTyVar subst) tpl_tvs) : ms) us rest
+      = ((item, map (lookupTyVar subst) tpl_tvs) : acc)
+      | otherwise
+      = acc
 
+
+    check_unifier :: [ClsInst] -> PotentialUnifiers
+    check_unifier [] = NoUnifiers
+    check_unifier (item@ClsInst { is_tvs = tpl_tvs, is_tys = tpl_tys }:items)
+      | not (instIsVisible vis_mods item)
+      = check_unifier items  -- See Note [Instance lookup and orphan instances]
+      | Just {} <- tcMatchTys tpl_tys tys = check_unifier items
         -- Does not match, so next check whether the things unify
         -- See Note [Overlapping instances]
         -- Ignore ones that are incoherent: Note [Incoherent instances]
       | isIncoherent item
-      = find ms us rest
+      = check_unifier items
 
       | otherwise
       = assertPpr (tys_tv_set `disjointVarSet` tpl_tv_set)
@@ -868,10 +914,12 @@ lookupInstEnv' ie vis_mods cls tys
           -- We consider MaybeApart to be a case where the instance might
           -- apply in the future. This covers an instance like C Int and
           -- a target like [W] C (F a), where F is a type family.
-            SurelyApart              -> find ms us        rest
+            SurelyApart              -> check_unifier items
               -- See Note [Infinitary substitution in lookup]
-            MaybeApart MARInfinite _ -> find ms us        rest
-            _                        -> find ms (item:us) rest
+            MaybeApart MARInfinite _ -> check_unifier items
+            _                        ->
+              OneOrMoreUnifiers (item: getPotentialUnifiers (check_unifier items))
+
       where
         tpl_tv_set = mkVarSet tpl_tvs
         tys_tv_set = tyCoVarsOfTypes tys
@@ -891,13 +939,12 @@ lookupInstEnv check_overlap_safe
                         , ie_visible = vis_mods })
               cls
               tys
-  = -- pprTrace "lookupInstEnv" (ppr cls <+> ppr tys $$ ppr home_ie) $
-    (final_matches, final_unifs, unsafe_overlapped)
+  = (final_matches, final_unifs, unsafe_overlapped)
   where
     (home_matches, home_unifs) = lookupInstEnv' home_ie vis_mods cls tys
     (pkg_matches,  pkg_unifs)  = lookupInstEnv' pkg_ie  vis_mods cls tys
     all_matches = home_matches ++ pkg_matches
-    all_unifs   = home_unifs   ++ pkg_unifs
+    all_unifs   = home_unifs   `mappend` pkg_unifs
     final_matches = pruneOverlappedMatches all_matches
         -- Even if the unifs is non-empty (an error situation)
         -- we still prune the matches, so that the error message isn't
@@ -911,7 +958,7 @@ lookupInstEnv check_overlap_safe
 
     -- If the selected match is incoherent, discard all unifiers
     final_unifs = case final_matches of
-                    (m:_) | isIncoherent (fst m) -> []
+                    (m:_) | isIncoherent (fst m) -> NoUnifiers
                     _                            -> all_unifs
 
     -- NOTE [Safe Haskell isSafeOverlap]