Cache the number of data cons in DataTyCon and SumTyCon

This is a follow-up after faf60e85 - Make tagForCon non-linear.
On the mailing list @simonpj suggested to solve the
linear behavior by caching the sizes.

Test Plan: ./validate

Reviewers: simonpj, simonmar, bgamari, austin

Reviewed By: simonpj

Subscribers: carter, goldfire, rwbarton, thomie, simonpj

Differential Revision: https://phabricator.haskell.org/D4131

6 files changed, 69 insertions, 82 deletions

diff --git a/compiler/codeGen/StgCmmClosure.hs b/compiler/codeGen/StgCmmClosure.hs
index 2501ec9cbd..1da1f707a2 100644
--- a/compiler/codeGen/StgCmmClosure.hs
+++ b/compiler/codeGen/StgCmmClosure.hs
@@ -361,18 +361,13 @@ type DynTag = Int       -- The tag on a *pointer*
 isSmallFamily :: DynFlags -> Int -> Bool
 isSmallFamily dflags fam_size = fam_size <= mAX_PTR_TAG dflags
 
--- | Faster version of isSmallFamily if you haven't computed the size yet.
-isSmallFamilyTyCon :: DynFlags -> TyCon -> Bool
-isSmallFamilyTyCon dflags tycon =
-  tyConFamilySizeAtMost tycon (mAX_PTR_TAG dflags)
-
 tagForCon :: DynFlags -> DataCon -> DynTag
 tagForCon dflags con
-  | isSmallFamilyTyCon dflags tycon = con_tag
-  | otherwise                       = 1
+  | isSmallFamily dflags fam_size = con_tag
+  | otherwise                     = 1
   where
     con_tag  = dataConTag con -- NB: 1-indexed
-    tycon = dataConTyCon con
+    fam_size = tyConFamilySize (dataConTyCon con)
 
 tagForArity :: DynFlags -> RepArity -> DynTag
 tagForArity dflags arity
diff --git a/compiler/iface/BuildTyCl.hs b/compiler/iface/BuildTyCl.hs
index 3d34e6f5fd..113ec12b63 100644
--- a/compiler/iface/BuildTyCl.hs
+++ b/compiler/iface/BuildTyCl.hs
@@ -9,7 +9,7 @@ module BuildTyCl (
         buildDataCon,
         buildPatSyn,
         TcMethInfo, buildClass,
-        mkNewTyConRhs, mkDataTyConRhs,
+        mkNewTyConRhs,
         newImplicitBinder, newTyConRepName
     ) where
 
@@ -41,19 +41,6 @@ import UniqSupply
 import Util
 import Outputable
 
-mkDataTyConRhs :: [DataCon] -> AlgTyConRhs
-mkDataTyConRhs cons
-  = DataTyCon {
-        data_cons = cons,
-        is_enum = not (null cons) && all is_enum_con cons
-                  -- See Note [Enumeration types] in TyCon
-    }
-  where
-    is_enum_con con
-       | (_univ_tvs, ex_tvs, eq_spec, theta, arg_tys, _res)
-           <- dataConFullSig con
-       = null ex_tvs && null eq_spec && null theta && null arg_tys
-
 
 mkNewTyConRhs :: Name -> TyCon -> DataCon -> TcRnIf m n AlgTyConRhs
 -- ^ Monadic because it makes a Name for the coercion TyCon
diff --git a/compiler/prelude/TysWiredIn.hs b/compiler/prelude/TysWiredIn.hs
index 2ee7e147a6..fda6b221f3 100644
--- a/compiler/prelude/TysWiredIn.hs
+++ b/compiler/prelude/TysWiredIn.hs
@@ -471,21 +471,17 @@ parrTyCon_RDR   = nameRdrName parrTyConName
 ************************************************************************
 -}
 
-pcNonEnumTyCon :: Name -> Maybe CType -> [TyVar] -> [DataCon] -> TyCon
--- Not an enumeration
-pcNonEnumTyCon = pcTyCon False
-
 -- This function assumes that the types it creates have all parameters at
 -- Representational role, and that there is no kind polymorphism.
-pcTyCon :: Bool -> Name -> Maybe CType -> [TyVar] -> [DataCon] -> TyCon
-pcTyCon is_enum name cType tyvars cons
+pcTyCon :: Name -> Maybe CType -> [TyVar] -> [DataCon] -> TyCon
+pcTyCon name cType tyvars cons
   = mkAlgTyCon name
                 (mkAnonTyConBinders tyvars)
                 liftedTypeKind
                 (map (const Representational) tyvars)
                 cType
                 []              -- No stupid theta
-                (DataTyCon cons is_enum)
+                (mkDataTyConRhs cons)
                 (VanillaAlgTyCon (mkPrelTyConRepName name))
                 False           -- Not in GADT syntax
 
@@ -569,16 +565,15 @@ pcSpecialDataCon dc_name arg_tys tycon rri
 typeNatKindCon, typeSymbolKindCon :: TyCon
 -- data Nat
 -- data Symbol
-typeNatKindCon    = pcTyCon False typeNatKindConName    Nothing [] []
-typeSymbolKindCon = pcTyCon False typeSymbolKindConName Nothing [] []
+typeNatKindCon    = pcTyCon typeNatKindConName    Nothing [] []
+typeSymbolKindCon = pcTyCon typeSymbolKindConName Nothing [] []
 
 typeNatKind, typeSymbolKind :: Kind
 typeNatKind    = mkTyConTy typeNatKindCon
 typeSymbolKind = mkTyConTy typeSymbolKindCon
 
 constraintKindTyCon :: TyCon
-constraintKindTyCon = pcTyCon False constraintKindTyConName
-                              Nothing [] []
+constraintKindTyCon = pcTyCon constraintKindTyConName Nothing [] []
 
 liftedTypeKind, constraintKind :: Kind
 liftedTypeKind   = tYPE liftedRepTy
@@ -1032,7 +1027,7 @@ heqSCSelId, coercibleSCSelId :: Id
     -- Kind: forall k1 k2. k1 -> k2 -> Constraint
     binders   = mkTemplateTyConBinders [liftedTypeKind, liftedTypeKind] (\ks -> ks)
     roles     = [Nominal, Nominal, Nominal, Nominal]
-    rhs       = DataTyCon { data_cons = [datacon], is_enum = False }
+    rhs       = mkDataTyConRhs [datacon]
 
     tvs       = binderVars binders
     sc_pred   = mkTyConApp eqPrimTyCon (mkTyVarTys tvs)
@@ -1050,7 +1045,7 @@ heqSCSelId, coercibleSCSelId :: Id
     -- Kind: forall k. k -> k -> Constraint
     binders   = mkTemplateTyConBinders [liftedTypeKind] (\[k] -> [k,k])
     roles     = [Nominal, Representational, Representational]
-    rhs       = DataTyCon { data_cons = [datacon], is_enum = False }
+    rhs       = mkDataTyConRhs [datacon]
 
     tvs@[k,a,b] = binderVars binders
     sc_pred     = mkTyConApp eqReprPrimTyCon (mkTyVarTys [k, k, a, b])
@@ -1092,7 +1087,7 @@ unicodeStarKindTyCon  = buildSynTyCon unicodeStarKindTyConName
                                        (tYPE liftedRepTy)
 
 runtimeRepTyCon :: TyCon
-runtimeRepTyCon = pcNonEnumTyCon runtimeRepTyConName Nothing []
+runtimeRepTyCon = pcTyCon runtimeRepTyConName Nothing []
                           (vecRepDataCon : tupleRepDataCon :
                            sumRepDataCon : runtimeRepSimpleDataCons)
 
@@ -1165,8 +1160,7 @@ liftedRepDataConTy, unliftedRepDataConTy,
   = map (mkTyConTy . promoteDataCon) runtimeRepSimpleDataCons
 
 vecCountTyCon :: TyCon
-vecCountTyCon = pcTyCon True vecCountTyConName Nothing []
-                        vecCountDataCons
+vecCountTyCon = pcTyCon vecCountTyConName Nothing [] vecCountDataCons
 
 -- See Note [Wiring in RuntimeRep]
 vecCountDataCons :: [DataCon]
@@ -1184,7 +1178,7 @@ vec2DataConTy, vec4DataConTy, vec8DataConTy, vec16DataConTy, vec32DataConTy,
   vec64DataConTy] = map (mkTyConTy . promoteDataCon) vecCountDataCons
 
 vecElemTyCon :: TyCon
-vecElemTyCon = pcTyCon True vecElemTyConName Nothing [] vecElemDataCons
+vecElemTyCon = pcTyCon vecElemTyConName Nothing [] vecElemDataCons
 
 -- See Note [Wiring in RuntimeRep]
 vecElemDataCons :: [DataCon]
@@ -1249,7 +1243,7 @@ charTy :: Type
 charTy = mkTyConTy charTyCon
 
 charTyCon :: TyCon
-charTyCon   = pcNonEnumTyCon charTyConName
+charTyCon   = pcTyCon charTyConName
                    (Just (CType NoSourceText Nothing
                                   (NoSourceText,fsLit "HsChar")))
                    [] [charDataCon]
@@ -1263,7 +1257,7 @@ intTy :: Type
 intTy = mkTyConTy intTyCon
 
 intTyCon :: TyCon
-intTyCon = pcNonEnumTyCon intTyConName
+intTyCon = pcTyCon intTyConName
                (Just (CType NoSourceText Nothing (NoSourceText,fsLit "HsInt")))
                  [] [intDataCon]
 intDataCon :: DataCon
@@ -1273,7 +1267,7 @@ wordTy :: Type
 wordTy = mkTyConTy wordTyCon
 
 wordTyCon :: TyCon
-wordTyCon = pcNonEnumTyCon wordTyConName
+wordTyCon = pcTyCon wordTyConName
             (Just (CType NoSourceText Nothing (NoSourceText, fsLit "HsWord")))
                [] [wordDataCon]
 wordDataCon :: DataCon
@@ -1283,10 +1277,10 @@ word8Ty :: Type
 word8Ty = mkTyConTy word8TyCon
 
 word8TyCon :: TyCon
-word8TyCon = pcNonEnumTyCon word8TyConName
-                      (Just (CType NoSourceText Nothing
-                             (NoSourceText, fsLit "HsWord8"))) []
-                      [word8DataCon]
+word8TyCon = pcTyCon word8TyConName
+                     (Just (CType NoSourceText Nothing
+                            (NoSourceText, fsLit "HsWord8"))) []
+                     [word8DataCon]
 word8DataCon :: DataCon
 word8DataCon = pcDataCon word8DataConName [] [wordPrimTy] word8TyCon
 
@@ -1294,7 +1288,7 @@ floatTy :: Type
 floatTy = mkTyConTy floatTyCon
 
 floatTyCon :: TyCon
-floatTyCon   = pcNonEnumTyCon floatTyConName
+floatTyCon   = pcTyCon floatTyConName
                       (Just (CType NoSourceText Nothing
                              (NoSourceText, fsLit "HsFloat"))) []
                       [floatDataCon]
@@ -1305,7 +1299,7 @@ doubleTy :: Type
 doubleTy = mkTyConTy doubleTyCon
 
 doubleTyCon :: TyCon
-doubleTyCon = pcNonEnumTyCon doubleTyConName
+doubleTyCon = pcTyCon doubleTyConName
                       (Just (CType NoSourceText Nothing
                              (NoSourceText,fsLit "HsDouble"))) []
                       [doubleDataCon]
@@ -1367,7 +1361,7 @@ boolTy :: Type
 boolTy = mkTyConTy boolTyCon
 
 boolTyCon :: TyCon
-boolTyCon = pcTyCon True boolTyConName
+boolTyCon = pcTyCon boolTyConName
                     (Just (CType NoSourceText Nothing
                            (NoSourceText, fsLit "HsBool")))
                     [] [falseDataCon, trueDataCon]
@@ -1381,7 +1375,7 @@ falseDataConId = dataConWorkId falseDataCon
 trueDataConId  = dataConWorkId trueDataCon
 
 orderingTyCon :: TyCon
-orderingTyCon = pcTyCon True orderingTyConName Nothing
+orderingTyCon = pcTyCon orderingTyConName Nothing
                         [] [ltDataCon, eqDataCon, gtDataCon]
 
 ltDataCon, eqDataCon, gtDataCon :: DataCon
@@ -1410,11 +1404,12 @@ mkListTy :: Type -> Type
 mkListTy ty = mkTyConApp listTyCon [ty]
 
 listTyCon :: TyCon
-listTyCon = buildAlgTyCon listTyConName alpha_tyvar [Representational]
-                          Nothing []
-                          (DataTyCon [nilDataCon, consDataCon] False )
-                          False
-                          (VanillaAlgTyCon $ mkPrelTyConRepName listTyConName)
+listTyCon =
+  buildAlgTyCon listTyConName alpha_tyvar [Representational]
+                Nothing []
+                (mkDataTyConRhs [nilDataCon, consDataCon])
+                False
+                (VanillaAlgTyCon $ mkPrelTyConRepName listTyConName)
 
 nilDataCon :: DataCon
 nilDataCon  = pcDataCon nilDataConName alpha_tyvar [] listTyCon
@@ -1431,7 +1426,7 @@ consDataCon = pcDataConWithFixity True {- Declared infix -}
 -- Wired-in type Maybe
 
 maybeTyCon :: TyCon
-maybeTyCon = pcTyCon False maybeTyConName Nothing alpha_tyvar
+maybeTyCon = pcTyCon maybeTyConName Nothing alpha_tyvar
                      [nothingDataCon, justDataCon]
 
 nothingDataCon :: DataCon
@@ -1537,7 +1532,7 @@ mkPArrTy ty  = mkTyConApp parrTyCon [ty]
 --     @PrelPArr@.
 --
 parrTyCon :: TyCon
-parrTyCon  = pcNonEnumTyCon parrTyConName Nothing alpha_tyvar [parrDataCon]
+parrTyCon  = pcTyCon parrTyConName Nothing alpha_tyvar [parrDataCon]
 
 parrDataCon :: DataCon
 parrDataCon  = pcDataCon
diff --git a/compiler/types/TyCon.hs b/compiler/types/TyCon.hs
index cf96a49403..333f52c7fb 100644
--- a/compiler/types/TyCon.hs
+++ b/compiler/types/TyCon.hs
@@ -34,6 +34,7 @@ module TyCon(
         mkLiftedPrimTyCon,
         mkTupleTyCon,
         mkSumTyCon,
+        mkDataTyConRhs,
         mkSynonymTyCon,
         mkFamilyTyCon,
         mkPromotedDataCon,
@@ -78,7 +79,7 @@ module TyCon(
         tyConDataCons, tyConDataCons_maybe,
         tyConSingleDataCon_maybe, tyConSingleDataCon,
         tyConSingleAlgDataCon_maybe,
-        tyConFamilySize, tyConFamilySizeAtMost,
+        tyConFamilySize,
         tyConStupidTheta,
         tyConArity,
         tyConRoles,
@@ -132,7 +133,7 @@ import {-# SOURCE #-} TysWiredIn ( runtimeRepTyCon, constraintKind
                                  , vecCountTyCon, vecElemTyCon, liftedTypeKind
                                  , mkFunKind, mkForAllKind )
 import {-# SOURCE #-} DataCon    ( DataCon, dataConExTyVars, dataConFieldLabels
-                                 , dataConTyCon )
+                                 , dataConTyCon, dataConFullSig )
 
 import Binary
 import Var
@@ -840,7 +841,8 @@ data AlgTyConRhs
                           --
                           -- INVARIANT: Kept in order of increasing 'DataCon'
                           -- tag (see the tag assignment in DataCon.mkDataCon)
-
+        data_cons_size :: Int,
+                          -- ^ Cached value: length data_cons
         is_enum :: Bool   -- ^ Cached value: is this an enumeration type?
                           --   See Note [Enumeration types]
     }
@@ -852,7 +854,8 @@ data AlgTyConRhs
     }
 
   | SumTyCon {
-        data_cons :: [DataCon]
+        data_cons :: [DataCon],
+        data_cons_size :: Int  -- ^ Cached value: length data_cons
     }
 
   -- | Information about those 'TyCon's derived from a @newtype@ declaration
@@ -886,6 +889,23 @@ data AlgTyConRhs
                              -- again check Trac #1072.
     }
 
+mkSumTyConRhs :: [DataCon] -> AlgTyConRhs
+mkSumTyConRhs data_cons = SumTyCon data_cons (length data_cons)
+
+mkDataTyConRhs :: [DataCon] -> AlgTyConRhs
+mkDataTyConRhs cons
+  = DataTyCon {
+        data_cons = cons,
+        data_cons_size = length cons,
+        is_enum = not (null cons) && all is_enum_con cons
+                  -- See Note [Enumeration types] in TyCon
+    }
+  where
+    is_enum_con con
+       | (_univ_tvs, ex_tvs, eq_spec, theta, arg_tys, _res)
+           <- dataConFullSig con
+       = null ex_tvs && null eq_spec && null theta && null arg_tys
+
 -- | Some promoted datacons signify extra info relevant to GHC. For example,
 -- the @IntRep@ constructor of @RuntimeRep@ corresponds to the 'IntRep'
 -- constructor of 'PrimRep'. This data structure allows us to store this
@@ -1491,7 +1511,7 @@ mkSumTyCon name binders res_kind arity tyvars cons parent
         tyConCType       = Nothing,
         algTcGadtSyntax  = False,
         algTcStupidTheta = [],
-        algTcRhs         = SumTyCon { data_cons = cons },
+        algTcRhs         = mkSumTyConRhs cons,
         algTcFields      = emptyDFsEnv,
         algTcParent      = parent
     }
@@ -2163,27 +2183,13 @@ tyConSingleAlgDataCon_maybe _        = Nothing
 tyConFamilySize  :: TyCon -> Int
 tyConFamilySize tc@(AlgTyCon { algTcRhs = rhs })
   = case rhs of
-      DataTyCon { data_cons = cons } -> length cons
+      DataTyCon { data_cons_size = size } -> size
       NewTyCon {}                    -> 1
       TupleTyCon {}                  -> 1
-      SumTyCon { data_cons = cons }  -> length cons
+      SumTyCon { data_cons_size = size }  -> size
       _                              -> pprPanic "tyConFamilySize 1" (ppr tc)
 tyConFamilySize tc = pprPanic "tyConFamilySize 2" (ppr tc)
 
--- | Determine if number of value constructors a 'TyCon' has is smaller
--- than n. Faster than tyConFamilySize tc <= n.
--- Panics if the 'TyCon' is not algebraic or a tuple
-tyConFamilySizeAtMost  :: TyCon -> Int -> Bool
-tyConFamilySizeAtMost tc@(AlgTyCon { algTcRhs = rhs }) n
-  = case rhs of
-      DataTyCon { data_cons = cons } -> lengthAtMost cons n
-      NewTyCon {}                    -> 1 <= n
-      TupleTyCon {}                  -> 1 <= n
-      SumTyCon { data_cons = cons }  -> lengthAtMost cons n
-      _                              -> pprPanic "tyConFamilySizeAtMost 1"
-                                          (ppr tc)
-tyConFamilySizeAtMost tc _ = pprPanic "tyConFamilySizeAtMost 2" (ppr tc)
-
 -- | Extract an 'AlgTyConRhs' with information about data constructors from an
 -- algebraic or tuple 'TyCon'. Panics for any other sort of 'TyCon'
 algTyConRhs :: TyCon -> AlgTyConRhs
diff --git a/compiler/vectorise/Vectorise/Generic/PData.hs b/compiler/vectorise/Vectorise/Generic/PData.hs
index 4862ce204b..353d6963b6 100644
--- a/compiler/vectorise/Vectorise/Generic/PData.hs
+++ b/compiler/vectorise/Vectorise/Generic/PData.hs
@@ -68,7 +68,7 @@ buildDataFamInst name' fam_tc vect_tc rhs
 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
 buildPDataTyConRhs orig_name vect_tc repr_tc repr
  = do data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
-      return $ DataTyCon { data_cons = [data_con], is_enum = False }
+      return $ mkDataTyConRhs [data_con]
 
 
 buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
@@ -113,7 +113,7 @@ buildPDatasTyCon orig_tc vect_tc repr
 buildPDatasTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
 buildPDatasTyConRhs orig_name vect_tc repr_tc repr
  = do data_con <- buildPDatasDataCon orig_name vect_tc repr_tc repr
-      return $ DataTyCon { data_cons = [data_con], is_enum = False }
+      return $ mkDataTyConRhs [data_con]
 
 
 buildPDatasDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
diff --git a/compiler/vectorise/Vectorise/Type/TyConDecl.hs b/compiler/vectorise/Vectorise/Type/TyConDecl.hs
index 75025501db..e71637981a 100644
--- a/compiler/vectorise/Vectorise/Type/TyConDecl.hs
+++ b/compiler/vectorise/Vectorise/Type/TyConDecl.hs
@@ -143,25 +143,29 @@ vectAlgTyConRhs tc (AbstractTyCon {})
   = do dflags <- getDynFlags
        cantVectorise dflags "Can't vectorise imported abstract type" (ppr tc)
 vectAlgTyConRhs _tc (DataTyCon { data_cons = data_cons
+                               , data_cons_size = data_cons_size
                                , is_enum   = is_enum
                                })
   = do { data_cons' <- mapM vectDataCon data_cons
        ; zipWithM_ defDataCon data_cons data_cons'
        ; return $ DataTyCon { data_cons = data_cons'
+                            , data_cons_size = data_cons_size
                             , is_enum   = is_enum
                             }
        }
 
 vectAlgTyConRhs tc (TupleTyCon { data_con = con })
-  = vectAlgTyConRhs tc (DataTyCon { data_cons = [con], is_enum = False })
+  = vectAlgTyConRhs tc (mkDataTyConRhs [con])
     -- I'm not certain this is what you want to do for tuples,
     -- but it's the behaviour we had before I refactored the
     -- representation of AlgTyConRhs to add tuples
 
-vectAlgTyConRhs tc (SumTyCon { data_cons = cons })
+vectAlgTyConRhs tc (SumTyCon { data_cons = cons
+                             , data_cons_size = data_cons_size })
   = -- FIXME (osa): I'm pretty sure this is broken.. TupleTyCon case is probably
     -- also broken when the tuple is unboxed.
     vectAlgTyConRhs tc (DataTyCon { data_cons = cons
+                                  , data_cons_size = data_cons_size
                                   , is_enum = all (((==) 0) . dataConRepArity) cons })
 
 vectAlgTyConRhs tc (NewTyCon {})