Define `Infinite` list and use where appropriate.

Also add perf test for infinite list fusion.

In particular, in `GHC.Core`, often we deal with infinite lists of roles. Also in a few locations we deal with infinite lists of names.

Thanks to simonpj for helping to write the Note [Fusion for `Infinite` lists].

6 files changed, 47 insertions, 56 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 1449e2331d..ad4e1b4ada 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -63,6 +63,7 @@ module GHC.Core.Coercion (
         splitForAllCo_ty_maybe, splitForAllCo_co_maybe,
 
         nthRole, tyConRolesX, tyConRolesRepresentational, setNominalRole_maybe,
+        tyConRoleListX, tyConRoleListRepresentational,
 
         pickLR,
 
@@ -154,6 +155,8 @@ import GHC.Builtin.Types.Prim
 import GHC.Data.List.SetOps
 import GHC.Data.Maybe
 import GHC.Types.Unique.FM
+import GHC.Data.List.Infinite (Infinite (..))
+import qualified GHC.Data.List.Infinite as Inf
 
 import GHC.Utils.Misc
 import GHC.Utils.Outputable
@@ -408,12 +411,10 @@ where co_rep1, co_rep2 are the coercions on the representations.
 --
 -- > decomposeCo 3 c [r1, r2, r3] = [nth r1 0 c, nth r2 1 c, nth r3 2 c]
 decomposeCo :: Arity -> Coercion
-            -> [Role]  -- the roles of the output coercions
-                       -- this must have at least as many
-                       -- entries as the Arity provided
+            -> Infinite Role  -- the roles of the output coercions
             -> [Coercion]
 decomposeCo arity co rs
-  = [mkNthCo r n co | (n,r) <- [0..(arity-1)] `zip` rs ]
+  = [mkNthCo r n co | (n,r) <- [0..(arity-1)] `zip` Inf.toList rs ]
            -- Remember, Nth is zero-indexed
 
 decomposeFunCo :: HasDebugCallStack
@@ -533,7 +534,7 @@ splitTyConAppCo_maybe :: Coercion -> Maybe (TyCon, [Coercion])
 splitTyConAppCo_maybe co
   | Just (ty, r) <- isReflCo_maybe co
   = do { (tc, tys) <- splitTyConApp_maybe ty
-       ; let args = zipWith mkReflCo (tyConRolesX r tc) tys
+       ; let args = zipWith mkReflCo (tyConRoleListX r tc) tys
        ; return (tc, args) }
 splitTyConAppCo_maybe (TyConAppCo _ tc cos) = Just (tc, cos)
 splitTyConAppCo_maybe (FunCo _ w arg res)     = Just (funTyCon, cos)
@@ -819,15 +820,14 @@ mkAppCo co arg
     -- Expand type synonyms; a TyConAppCo can't have a type synonym (#9102)
   = mkTyConAppCo r tc (zip_roles (tyConRolesX r tc) tys)
   where
-    zip_roles (r1:_)  []            = [downgradeRole r1 Nominal arg]
-    zip_roles (r1:rs) (ty1:tys)     = mkReflCo r1 ty1 : zip_roles rs tys
-    zip_roles _       _             = panic "zip_roles" -- but the roles are infinite...
+    zip_roles (Inf r1 _)  []            = [downgradeRole r1 Nominal arg]
+    zip_roles (Inf r1 rs) (ty1:tys)     = mkReflCo r1 ty1 : zip_roles rs tys
 
 mkAppCo (TyConAppCo r tc args) arg
   = case r of
       Nominal          -> mkTyConAppCo Nominal tc (args ++ [arg])
       Representational -> mkTyConAppCo Representational tc (args ++ [arg'])
-        where new_role = (tyConRolesRepresentational tc) !! (length args)
+        where new_role = tyConRolesRepresentational tc Inf.!! length args
               arg'     = downgradeRole new_role Nominal arg
       Phantom          -> mkTyConAppCo Phantom tc (args ++ [toPhantomCo arg])
 mkAppCo co arg = AppCo co  arg
@@ -1153,10 +1153,7 @@ mkNthCo r n co
       , tc1 == tc2
       = let len1 = length tys1
             len2 = length tys2
-            good_role = case coercionRole co of
-                          Nominal -> r == Nominal
-                          Representational -> r == (tyConRolesRepresentational tc1 !! n)
-                          Phantom -> r == Phantom
+            good_role = r == nthRole (coercionRole co) tc1 n
         in len1 == len2 && n < len1 && good_role
 
       | otherwise
@@ -1349,7 +1346,7 @@ setNominalRole_maybe r co
     setNominalRole_maybe_helper co@(Refl _) = Just co
     setNominalRole_maybe_helper (GRefl _ ty co) = Just $ GRefl Nominal ty co
     setNominalRole_maybe_helper (TyConAppCo Representational tc cos)
-      = do { cos' <- zipWithM setNominalRole_maybe (tyConRolesX Representational tc) cos
+      = do { cos' <- zipWithM setNominalRole_maybe (tyConRoleListX Representational tc) cos
            ; return $ TyConAppCo Nominal tc cos' }
     setNominalRole_maybe_helper (FunCo Representational w co1 co2)
       = do { co1' <- setNominalRole_maybe Representational co1
@@ -1393,27 +1390,33 @@ toPhantomCo co
 
 -- Convert args to a TyConAppCo Nominal to the same TyConAppCo Representational
 applyRoles :: TyCon -> [Coercion] -> [Coercion]
-applyRoles tc cos
-  = zipWith (\r -> downgradeRole r Nominal) (tyConRolesRepresentational tc) cos
+applyRoles = zipWith (`downgradeRole` Nominal) . tyConRoleListRepresentational
 
 -- the Role parameter is the Role of the TyConAppCo
 -- defined here because this is intimately concerned with the implementation
 -- of TyConAppCo
 -- Always returns an infinite list (with a infinite tail of Nominal)
-tyConRolesX :: Role -> TyCon -> [Role]
+tyConRolesX :: Role -> TyCon -> Infinite Role
 tyConRolesX Representational tc = tyConRolesRepresentational tc
-tyConRolesX role             _  = repeat role
+tyConRolesX role             _  = Inf.repeat role
+
+tyConRoleListX :: Role -> TyCon -> [Role]
+tyConRoleListX role = Inf.toList . tyConRolesX role
+
+-- Returns the roles of the parameters of a tycon, with an infinite tail
+-- of Nominal
+tyConRolesRepresentational :: TyCon -> Infinite Role
+tyConRolesRepresentational tc = tyConRoles tc Inf.++ Inf.repeat Nominal
 
 -- Returns the roles of the parameters of a tycon, with an infinite tail
 -- of Nominal
-tyConRolesRepresentational :: TyCon -> [Role]
-tyConRolesRepresentational tc = tyConRoles tc ++ repeat Nominal
+tyConRoleListRepresentational :: TyCon -> [Role]
+tyConRoleListRepresentational = Inf.toList . tyConRolesRepresentational
 
 nthRole :: Role -> TyCon -> Int -> Role
 nthRole Nominal _ _ = Nominal
 nthRole Phantom _ _ = Phantom
-nthRole Representational tc n
-  = (tyConRolesRepresentational tc) `getNth` n
+nthRole Representational tc n = tyConRolesRepresentational tc Inf.!! n
 
 ltRole :: Role -> Role -> Bool
 -- Is one role "less" than another?
@@ -2034,7 +2037,7 @@ ty_co_subst !lc role ty
     go r (TyVarTy tv)      = expectJust "ty_co_subst bad roles" $
                              liftCoSubstTyVar lc r tv
     go r (AppTy ty1 ty2)   = mkAppCo (go r ty1) (go Nominal ty2)
-    go r (TyConApp tc tys) = mkTyConAppCo r tc (zipWith go (tyConRolesX r tc) tys)
+    go r (TyConApp tc tys) = mkTyConAppCo r tc (zipWith go (tyConRoleListX r tc) tys)
     go r (FunTy _ w ty1 ty2) = mkFunCo r (go Nominal w) (go r ty1) (go r ty2)
     go r t@(ForAllTy (Bndr v _) ty)
        = let (lc', v', h) = liftCoSubstVarBndr lc v
diff --git a/compiler/GHC/Core/Coercion/Opt.hs b/compiler/GHC/Core/Coercion/Opt.hs
index 6fa8fc1273..d061d795a7 100644
--- a/compiler/GHC/Core/Coercion/Opt.hs
+++ b/compiler/GHC/Core/Coercion/Opt.hs
@@ -245,7 +245,7 @@ opt_co4 env sym rep r g@(TyConAppCo _r tc cos)
       (True, Nominal) ->
         mkTyConAppCo Representational tc
                      (zipWith3 (opt_co3 env sym)
-                               (map Just (tyConRolesRepresentational tc))
+                               (map Just (tyConRoleListRepresentational tc))
                                (repeat Nominal)
                                cos)
       (False, Nominal) ->
@@ -254,7 +254,7 @@ opt_co4 env sym rep r g@(TyConAppCo _r tc cos)
                       -- must use opt_co2 here, because some roles may be P
                       -- See Note [Optimising coercion optimisation]
         mkTyConAppCo r tc (zipWith (opt_co2 env sym)
-                                   (tyConRolesRepresentational tc)  -- the current roles
+                                   (tyConRoleListRepresentational tc)  -- the current roles
                                    cos)
       (_, Phantom) -> pprPanic "opt_co4 sees a phantom!" (ppr g)
 
@@ -546,7 +546,7 @@ opt_univ env sym prov role oty1 oty2
   , equalLength tys1 tys2 -- see Note [Differing kinds]
       -- NB: prov must not be the two interesting ones (ProofIrrel & Phantom);
       -- Phantom is already taken care of, and ProofIrrel doesn't relate tyconapps
-  = let roles    = tyConRolesX role tc1
+  = let roles    = tyConRoleListX role tc1
         arg_cos  = zipWith3 (mkUnivCo prov') roles tys1 tys2
         arg_cos' = zipWith (opt_co4 env sym False) roles arg_cos
     in
diff --git a/compiler/GHC/Core/FamInstEnv.hs b/compiler/GHC/Core/FamInstEnv.hs
index e955e5befd..5ecb83d4a6 100644
--- a/compiler/GHC/Core/FamInstEnv.hs
+++ b/compiler/GHC/Core/FamInstEnv.hs
@@ -63,6 +63,8 @@ import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Utils.Panic.Plain
 import GHC.Data.Bag
+import GHC.Data.List.Infinite (Infinite (..))
+import qualified GHC.Data.List.Infinite as Inf
 
 {-
 ************************************************************************
@@ -1477,7 +1479,7 @@ normalise_type ty
                Nothing ->
                  do { ArgsReductions redns res_co
                         <- normalise_args (typeKind nfun)
-                                          (repeat Nominal)
+                                          (Inf.repeat Nominal)
                                           arg_tys
                     ; role <- getRole
                     ; return $
@@ -1486,7 +1488,7 @@ normalise_type ty
                           (mkSymMCo res_co) } }
 
 normalise_args :: Kind    -- of the function
-               -> [Role]  -- roles at which to normalise args
+               -> Infinite Role  -- roles at which to normalise args
                -> [Type]  -- args
                -> NormM ArgsReductions
 -- returns ArgsReductions (Reductions cos xis) res_co,
@@ -1496,7 +1498,7 @@ normalise_args :: Kind    -- of the function
 -- but the resulting application *will* be well-kinded
 -- cf. GHC.Tc.Solver.Rewrite.rewrite_args_slow
 normalise_args fun_ki roles args
-  = do { normed_args <- zipWithM normalise1 roles args
+  = do { normed_args <- zipWithM normalise1 (Inf.toList roles) args
        ; return $ simplifyArgsWorker ki_binders inner_ki fvs roles normed_args }
   where
     (ki_binders, inner_ki) = splitPiTys fun_ki
diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index 086a727095..6c285db819 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -2177,7 +2177,7 @@ lintCoercion co@(TyConAppCo r tc cos)
        ; let (co_kinds, co_roles) = unzip (map coercionKindRole cos')
        ; lint_co_app co (tyConKind tc) (map pFst co_kinds)
        ; lint_co_app co (tyConKind tc) (map pSnd co_kinds)
-       ; zipWithM_ (lintRole co) (tyConRolesX r tc) co_roles
+       ; zipWithM_ (lintRole co) (tyConRoleListX r tc) co_roles
        ; return (TyConAppCo r tc cos') }
 
 lintCoercion co@(AppCo co1 co2)
diff --git a/compiler/GHC/Core/Reduction.hs b/compiler/GHC/Core/Reduction.hs
index f15b335fd7..f97b9517b6 100644
--- a/compiler/GHC/Core/Reduction.hs
+++ b/compiler/GHC/Core/Reduction.hs
@@ -35,6 +35,8 @@ import GHC.Core.TyCon      ( TyCon )
 import GHC.Core.Type
 
 import GHC.Data.Pair       ( Pair(Pair) )
+import GHC.Data.List.Infinite ( Infinite (..) )
+import qualified GHC.Data.List.Infinite as Inf
 
 import GHC.Types.Var       ( setTyVarKind )
 import GHC.Types.Var.Env   ( mkInScopeSet )
@@ -42,7 +44,7 @@ import GHC.Types.Var.Set   ( TyCoVarSet )
 
 import GHC.Utils.Misc      ( HasDebugCallStack, equalLength )
 import GHC.Utils.Outputable
-import GHC.Utils.Panic     ( assertPpr, panic )
+import GHC.Utils.Panic     ( assertPpr )
 
 {-
 %************************************************************************
@@ -788,7 +790,7 @@ simplifyArgsWorker :: HasDebugCallStack
                        -- the binders & result kind (not a Π-type) of the function applied to the args
                        -- list of binders can be shorter or longer than the list of args
                    -> TyCoVarSet   -- free vars of the args
-                   -> [Role]       -- list of roles, r
+                   -> Infinite Role-- list of roles, r
                    -> [Reduction]  -- rewritten type arguments, arg_i
                                    -- each comes with the coercion used to rewrite it,
                                    -- arg_co_i :: ty_i ~ arg_i
@@ -809,11 +811,11 @@ simplifyArgsWorker orig_ki_binders orig_inner_ki orig_fvs
   where
     orig_lc = emptyLiftingContext $ mkInScopeSet orig_fvs
 
-    go :: LiftingContext  -- mapping from tyvars to rewriting coercions
-       -> [TyCoBinder]    -- Unsubsted binders of function's kind
-       -> Kind        -- Unsubsted result kind of function (not a Pi-type)
-       -> [Role]      -- Roles at which to rewrite these ...
-       -> [Reduction] -- rewritten arguments, with their rewriting coercions
+    go :: LiftingContext -- mapping from tyvars to rewriting coercions
+       -> [TyCoBinder]   -- Unsubsted binders of function's kind
+       -> Kind           -- Unsubsted result kind of function (not a Pi-type)
+       -> Infinite Role  -- Roles at which to rewrite these ...
+       -> [Reduction]    -- rewritten arguments, with their rewriting coercions
        -> ArgsReductions
     go !lc binders inner_ki _ []
         -- The !lc makes the function strict in the lifting context
@@ -826,7 +828,7 @@ simplifyArgsWorker orig_ki_binders orig_inner_ki orig_fvs
         kind_co | noFreeVarsOfType final_kind = MRefl
                 | otherwise                   = MCo $ liftCoSubst Nominal lc final_kind
 
-    go lc (binder:binders) inner_ki (role:roles) (arg_redn:arg_redns)
+    go lc (binder:binders) inner_ki (Inf role roles) (arg_redn:arg_redns)
       =  -- We rewrite an argument ty with arg_redn = Reduction arg_co arg
          -- By Note [Rewriting] in GHC.Tc.Solver.Rewrite invariant (F2),
          -- tcTypeKind(ty) = tcTypeKind(arg).
@@ -859,7 +861,7 @@ simplifyArgsWorker orig_ki_binders orig_inner_ki orig_fvs
             (arg_cos, res_co)     = decomposePiCos co1 co1_kind unrewritten_tys
             casted_args           = assertPpr (equalLength arg_redns arg_cos)
                                               (ppr arg_redns $$ ppr arg_cos)
-                                  $ zipWith3 mkCoherenceRightRedn roles arg_redns arg_cos
+                                  $ zipWith3 mkCoherenceRightRedn (Inf.toList roles) arg_redns arg_cos
                -- In general decomposePiCos can return fewer cos than tys,
                -- but not here; because we're well typed, there will be enough
                -- binders. Note that decomposePiCos does substitutions, so even
@@ -874,19 +876,3 @@ simplifyArgsWorker orig_ki_binders orig_inner_ki orig_fvs
               = go zapped_lc bndrs new_inner roles casted_args
         in
           ArgsReductions redns_out (res_co `mkTransMCoR` res_co_out)
-
-    go _ _ _ _ _ = panic
-        "simplifyArgsWorker wandered into deeper water than usual"
-           -- This debug information is commented out because leaving it in
-           -- causes a ~2% increase in allocations in T9872d.
-           -- That's independent of the analogous case in rewrite_args_fast
-           -- in GHC.Tc.Solver.Rewrite:
-           -- each of these causes a 2% increase on its own, so commenting them
-           -- both out gives a 4% decrease in T9872d.
-           {-
-
-             (vcat [ppr orig_binders,
-                    ppr orig_inner_ki,
-                    ppr (take 10 orig_roles), -- often infinite!
-                    ppr orig_tys])
-           -}
diff --git a/compiler/GHC/Core/Unify.hs b/compiler/GHC/Core/Unify.hs
index 188d5ff32f..596fef6b6f 100644
--- a/compiler/GHC/Core/Unify.hs
+++ b/compiler/GHC/Core/Unify.hs
@@ -1742,7 +1742,7 @@ pushRefl co =
       ->  Just (TyConAppCo r funTyCon [ multToCo w, mkReflCo r rep1, mkReflCo r rep2
                                        , mkReflCo r ty1,  mkReflCo r ty2 ])
     Just (TyConApp tc tys, r)
-      -> Just (TyConAppCo r tc (zipWith mkReflCo (tyConRolesX r tc) tys))
+      -> Just (TyConAppCo r tc (zipWith mkReflCo (tyConRoleListX r tc) tys))
     Just (ForAllTy (Bndr tv _) ty, r)
       -> Just (ForAllCo tv (mkNomReflCo (varType tv)) (mkReflCo r ty))
     -- NB: NoRefl variant. Otherwise, we get a loop!