Speed up core size and core stats

Summary:
When calculating core size and core stats, we previously
calculated sizes/stats for sub-parts and then added them.
It should be faster to thread an accumulator through.

Reviewers: austin, bgamari

Subscribers: rwbarton, thomie

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

2 files changed, 91 insertions, 64 deletions

diff --git a/compiler/coreSyn/CoreStats.hs b/compiler/coreSyn/CoreStats.hs
index cb73d147a8..5827f163f5 100644
--- a/compiler/coreSyn/CoreStats.hs
+++ b/compiler/coreSyn/CoreStats.hs
@@ -27,6 +27,7 @@ data CoreStats = CS { cs_tm :: !Int    -- Terms
                     , cs_vb :: !Int    -- Local value bindings
                     , cs_jb :: !Int }  -- Local join bindings
 
+type CST = CoreStats -> CoreStats
 
 instance Outputable CoreStats where
  ppr (CS { cs_tm = i1, cs_ty = i2, cs_co = i3, cs_vb = i4, cs_jb = i5 })
@@ -42,59 +43,69 @@ plusCS (CS { cs_tm = p1, cs_ty = q1, cs_co = r1, cs_vb = v1, cs_jb = j1 })
   = CS { cs_tm = p1+p2, cs_ty = q1+q2, cs_co = r1+r2, cs_vb = v1+v2
        , cs_jb = j1+j2 }
 
+plusCST :: CST -> CST -> CST
+plusCST f g acc = g $! f acc
+
+zeroCST, oneTMT :: CST
+zeroCST = plusCS zeroCS
+oneTMT = plusCS oneTM
+
 zeroCS, oneTM :: CoreStats
 zeroCS = CS { cs_tm = 0, cs_ty = 0, cs_co = 0, cs_vb = 0, cs_jb = 0 }
 oneTM  = zeroCS { cs_tm = 1 }
 
-sumCS :: (a -> CoreStats) -> [a] -> CoreStats
-sumCS f = foldl' (\s a -> plusCS s (f a)) zeroCS
+sumCST :: (a -> CST) -> [a] -> CST
+sumCST f = foldr (\a r s -> r $! f a $! s) zeroCST
 
 coreBindsStats :: [CoreBind] -> CoreStats
-coreBindsStats = sumCS (bindStats TopLevel)
+coreBindsStats binds = sumCST (bindStats TopLevel) binds zeroCS
 
-bindStats :: TopLevelFlag -> CoreBind -> CoreStats
+bindStats :: TopLevelFlag -> CoreBind -> CST
 bindStats top_lvl (NonRec v r) = bindingStats top_lvl v r
-bindStats top_lvl (Rec prs)    = sumCS (\(v,r) -> bindingStats top_lvl v r) prs
+bindStats top_lvl (Rec prs)    = sumCST (\(v,r) -> bindingStats top_lvl v r) prs
 
-bindingStats :: TopLevelFlag -> Var -> CoreExpr -> CoreStats
-bindingStats top_lvl v r = letBndrStats top_lvl v `plusCS` exprStats r
+bindingStats :: TopLevelFlag -> Var -> CoreExpr -> CST
+bindingStats top_lvl v r = letBndrStats top_lvl v `plusCST` exprStatsT r
 
-bndrStats :: Var -> CoreStats
-bndrStats v = oneTM `plusCS` tyStats (varType v)
+bndrStats :: Var -> CST
+bndrStats v = oneTMT `plusCST` tyStats (varType v)
 
-letBndrStats :: TopLevelFlag -> Var -> CoreStats
+letBndrStats :: TopLevelFlag -> Var -> CST
 letBndrStats top_lvl v
   | isTyVar v || isTopLevel top_lvl = bndrStats v
-  | isJoinId v = oneTM { cs_jb = 1 } `plusCS` ty_stats
-  | otherwise  = oneTM { cs_vb = 1 } `plusCS` ty_stats
+  | isJoinId v = plusCS (oneTM { cs_jb = 1 }) `plusCST` ty_stats
+  | otherwise  = plusCS (oneTM { cs_vb = 1 }) `plusCST` ty_stats
   where
     ty_stats = tyStats (varType v)
 
 exprStats :: CoreExpr -> CoreStats
-exprStats (Var {})        = oneTM
-exprStats (Lit {})        = oneTM
-exprStats (Type t)        = tyStats t
-exprStats (Coercion c)    = coStats c
-exprStats (App f a)       = exprStats f `plusCS` exprStats a
-exprStats (Lam b e)       = bndrStats b `plusCS` exprStats e
-exprStats (Let b e)       = bindStats NotTopLevel b `plusCS` exprStats e
-exprStats (Case e b _ as) = exprStats e `plusCS` bndrStats b
-                                        `plusCS` sumCS altStats as
-exprStats (Cast e co)     = coStats co `plusCS` exprStats e
-exprStats (Tick _ e)      = exprStats e
-
-altStats :: CoreAlt -> CoreStats
-altStats (_, bs, r) = altBndrStats bs `plusCS` exprStats r
-
-altBndrStats :: [Var] -> CoreStats
+exprStats e = exprStatsT e zeroCS
+
+exprStatsT :: CoreExpr -> CST
+exprStatsT (Var {})        = oneTMT
+exprStatsT (Lit {})        = oneTMT
+exprStatsT (Type t)        = tyStats t
+exprStatsT (Coercion c)    = coStats c
+exprStatsT (App f a)       = exprStatsT f `plusCST` exprStatsT a
+exprStatsT (Lam b e)       = bndrStats b `plusCST` exprStatsT e
+exprStatsT (Let b e)       = bindStats NotTopLevel b `plusCST` exprStatsT e
+exprStatsT (Case e b _ as) = exprStatsT e `plusCST` bndrStats b
+                                        `plusCST` sumCST altStats as
+exprStatsT (Cast e co)     = coStats co `plusCST` exprStatsT e
+exprStatsT (Tick _ e)      = exprStatsT e
+
+altStats :: CoreAlt -> CST
+altStats (_, bs, r) = altBndrStats bs `plusCST` exprStatsT r
+
+altBndrStats :: [Var] -> CST
 -- Charge one for the alternative, not for each binder
-altBndrStats vs = oneTM `plusCS` sumCS (tyStats . varType) vs
+altBndrStats vs = oneTMT `plusCST` sumCST (tyStats . varType) vs
 
-tyStats :: Type -> CoreStats
-tyStats ty = zeroCS { cs_ty = typeSize ty }
+tyStats :: Type -> CST
+tyStats ty = plusCS $ zeroCS { cs_ty = typeSize ty }
 
-coStats :: Coercion -> CoreStats
-coStats co = zeroCS { cs_co = coercionSize co }
+coStats :: Coercion -> CST
+coStats co = plusCS $ zeroCS { cs_co = coercionSize co }
 
 coreBindsSize :: [CoreBind] -> Int
 -- We use coreBindStats for user printout
diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index 5ac63e5b04..e4fe0e8916 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -128,7 +128,7 @@ module TyCoRep (
         tidyTyVarBinder, tidyTyVarBinders,
 
         -- * Sizes
-        typeSize, coercionSize, provSize
+        typeSize, coercionSize, typeSizePlus, coercionSizePlus, provSize
     ) where
 
 #include "HsVersions.h"
@@ -2824,37 +2824,53 @@ tidyCos env = map (tidyCo env)
 --     function is used only in reporting, not decision-making.
 
 typeSize :: Type -> Int
-typeSize (LitTy {})                 = 1
-typeSize (TyVarTy {})               = 1
-typeSize (AppTy t1 t2)              = typeSize t1 + typeSize t2
-typeSize (FunTy t1 t2)              = typeSize t1 + typeSize t2
-typeSize (ForAllTy (TvBndr tv _) t) = typeSize (tyVarKind tv) + typeSize t
-typeSize (TyConApp _ ts)            = 1 + sum (map typeSize ts)
-typeSize (CastTy ty co)             = typeSize ty + coercionSize co
-typeSize (CoercionTy co)            = coercionSize co
+typeSize t = typeSizePlus t 0
 
 coercionSize :: Coercion -> Int
-coercionSize (Refl _ ty)         = typeSize ty
-coercionSize (TyConAppCo _ _ args) = 1 + sum (map coercionSize args)
-coercionSize (AppCo co arg)      = coercionSize co + coercionSize arg
-coercionSize (ForAllCo _ h co)   = 1 + coercionSize co + coercionSize h
-coercionSize (FunCo _ co1 co2)   = 1 + coercionSize co1 + coercionSize co2
-coercionSize (CoVarCo _)         = 1
-coercionSize (AxiomInstCo _ _ args) = 1 + sum (map coercionSize args)
-coercionSize (UnivCo p _ t1 t2)  = 1 + provSize p + typeSize t1 + typeSize t2
-coercionSize (SymCo co)          = 1 + coercionSize co
-coercionSize (TransCo co1 co2)   = 1 + coercionSize co1 + coercionSize co2
-coercionSize (NthCo _ co)        = 1 + coercionSize co
-coercionSize (LRCo  _ co)        = 1 + coercionSize co
-coercionSize (InstCo co arg)     = 1 + coercionSize co + coercionSize arg
-coercionSize (CoherenceCo c1 c2) = 1 + coercionSize c1 + coercionSize c2
-coercionSize (KindCo co)         = 1 + coercionSize co
-coercionSize (SubCo co)          = 1 + coercionSize co
-coercionSize (AxiomRuleCo _ cs)  = 1 + sum (map coercionSize cs)
+coercionSize t = coercionSizePlus t 0
+
+sizerPlus :: (Int -> Int) -> (Int -> Int) -> Int -> Int
+sizerPlus f g x = g $! f $! x
+{-# INLINE sizerPlus #-}
+
+typeSizePlus :: Type -> Int -> Int
+typeSizePlus (LitTy {})                 = (+ 1)
+typeSizePlus (TyVarTy {})               = (+ 1)
+typeSizePlus (AppTy t1 t2)              = typeSizePlus t1 `sizerPlus` typeSizePlus t2
+typeSizePlus (FunTy t1 t2)              = typeSizePlus t1 `sizerPlus` typeSizePlus t2
+typeSizePlus (ForAllTy (TvBndr tv _) t) = typeSizePlus (tyVarKind tv) `sizerPlus` typeSizePlus t
+typeSizePlus (TyConApp _ ts)            = \acc0 ->
+                        foldl' (\acc arg -> typeSizePlus arg acc) (acc0 + 1) ts
+typeSizePlus (CastTy ty co)             = typeSizePlus ty `sizerPlus` coercionSizePlus co
+typeSizePlus (CoercionTy co)            = coercionSizePlus co
+
+coercionSizePlus (Refl _ ty)         = typeSizePlus ty
+coercionSizePlus (TyConAppCo _ _ args) = \acc0 ->
+                        foldl' (\acc arg -> coercionSizePlus arg acc) (acc0 + 1) args
+coercionSizePlus (AppCo co arg)      = coercionSizePlus co `sizerPlus` coercionSizePlus arg
+coercionSizePlus (ForAllCo _ h co)   = (+ 1) `sizerPlus` coercionSizePlus co `sizerPlus` coercionSizePlus h
+coercionSizePlus (FunCo _ co1 co2)   = (+ 1) `sizerPlus` coercionSizePlus co1 `sizerPlus` coercionSizePlus co2
+coercionSizePlus (CoVarCo _)         = (+ 1)
+coercionSizePlus (AxiomInstCo _ _ args) = \acc0 ->
+                        foldl' (\acc arg -> coercionSizePlus arg acc) (acc0 + 1) args
+coercionSizePlus (UnivCo p _ t1 t2)  = (+ 1) `sizerPlus` provSizePlus p `sizerPlus` typeSizePlus t1 `sizerPlus` typeSizePlus t2
+coercionSizePlus (SymCo co)          = (+ 1) `sizerPlus` coercionSizePlus co
+coercionSizePlus (TransCo co1 co2)   = (+ 1) `sizerPlus` coercionSizePlus co1 `sizerPlus` coercionSizePlus co2
+coercionSizePlus (NthCo _ co)        = (+ 1) `sizerPlus` coercionSizePlus co
+coercionSizePlus (LRCo  _ co)        = (+ 1) `sizerPlus` coercionSizePlus co
+coercionSizePlus (InstCo co arg)     = (+ 1) `sizerPlus`  coercionSizePlus co `sizerPlus` coercionSizePlus arg
+coercionSizePlus (CoherenceCo c1 c2) = (+ 1) `sizerPlus` coercionSizePlus c1 `sizerPlus` coercionSizePlus c2
+coercionSizePlus (KindCo co)         = (+ 1) `sizerPlus` coercionSizePlus co
+coercionSizePlus (SubCo co)          = (+ 1) `sizerPlus` coercionSizePlus co
+coercionSizePlus (AxiomRuleCo _ cs)  = \acc0 ->
+                        foldl' (\acc arg -> coercionSizePlus arg acc) (acc0 + 1) cs
 
 provSize :: UnivCoProvenance -> Int
-provSize UnsafeCoerceProv    = 1
-provSize (PhantomProv co)    = 1 + coercionSize co
-provSize (ProofIrrelProv co) = 1 + coercionSize co
-provSize (PluginProv _)      = 1
-provSize (HoleProv h)        = pprPanic "provSize hits a hole" (ppr h)
+provSize prov = provSizePlus prov 0
+
+provSizePlus :: UnivCoProvenance -> Int -> Int
+provSizePlus UnsafeCoerceProv    = (+ 1)
+provSizePlus (PhantomProv co)    = (+ 1) `sizerPlus` coercionSizePlus co
+provSizePlus (ProofIrrelProv co) = (+ 1) `sizerPlus` coercionSizePlus co
+provSizePlus (PluginProv _)      = (+ 1)
+provSizePlus (HoleProv h)        = \a -> a `seq` pprPanic "provSizePlus hits a hole" (ppr h)