{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-2015
-}

-- | Functions to computing the statistics reflective of the "size"
-- of a Core expression
module CoreStats (
        -- * Expression and bindings size
        coreBindsSize, exprSize,
        CoreStats(..), coreBindsStats, exprStats,
    ) where

import GhcPrelude

import BasicTypes
import CoreSyn
import Outputable
import Coercion
import Var
import Type (Type, typeSize)
import Id (isJoinId)

data CoreStats = CS { CoreStats -> Int
cs_tm :: !Int    -- Terms
                    , CoreStats -> Int
cs_ty :: !Int    -- Types
                    , CoreStats -> Int
cs_co :: !Int    -- Coercions
                    , CoreStats -> Int
cs_vb :: !Int    -- Local value bindings
                    , CoreStats -> Int
cs_jb :: !Int }  -- Local join bindings


instance Outputable CoreStats where
 ppr :: CoreStats -> SDoc
ppr (CS { cs_tm :: CoreStats -> Int
cs_tm = Int
i1, cs_ty :: CoreStats -> Int
cs_ty = Int
i2, cs_co :: CoreStats -> Int
cs_co = Int
i3, cs_vb :: CoreStats -> Int
cs_vb = Int
i4, cs_jb :: CoreStats -> Int
cs_jb = Int
i5 })
   = SDoc -> SDoc
braces ([SDoc] -> SDoc
sep [String -> SDoc
text "terms:"     SDoc -> SDoc -> SDoc
<+> Int -> SDoc
forall a. Integral a => a -> SDoc
intWithCommas Int
i1 SDoc -> SDoc -> SDoc
<> SDoc
comma,
                  String -> SDoc
text "types:"     SDoc -> SDoc -> SDoc
<+> Int -> SDoc
forall a. Integral a => a -> SDoc
intWithCommas Int
i2 SDoc -> SDoc -> SDoc
<> SDoc
comma,
                  String -> SDoc
text "coercions:" SDoc -> SDoc -> SDoc
<+> Int -> SDoc
forall a. Integral a => a -> SDoc
intWithCommas Int
i3 SDoc -> SDoc -> SDoc
<> SDoc
comma,
                  String -> SDoc
text "joins:"     SDoc -> SDoc -> SDoc
<+> Int -> SDoc
forall a. Integral a => a -> SDoc
intWithCommas Int
i5 SDoc -> SDoc -> SDoc
<> Char -> SDoc
char '/' SDoc -> SDoc -> SDoc
<>
                                        Int -> SDoc
forall a. Integral a => a -> SDoc
intWithCommas (Int
i4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i5) ])

plusCS :: CoreStats -> CoreStats -> CoreStats
plusCS :: CoreStats -> CoreStats -> CoreStats
plusCS (CS { cs_tm :: CoreStats -> Int
cs_tm = Int
p1, cs_ty :: CoreStats -> Int
cs_ty = Int
q1, cs_co :: CoreStats -> Int
cs_co = Int
r1, cs_vb :: CoreStats -> Int
cs_vb = Int
v1, cs_jb :: CoreStats -> Int
cs_jb = Int
j1 })
       (CS { cs_tm :: CoreStats -> Int
cs_tm = Int
p2, cs_ty :: CoreStats -> Int
cs_ty = Int
q2, cs_co :: CoreStats -> Int
cs_co = Int
r2, cs_vb :: CoreStats -> Int
cs_vb = Int
v2, cs_jb :: CoreStats -> Int
cs_jb = Int
j2 })
  = $WCS :: Int -> Int -> Int -> Int -> Int -> CoreStats
CS { cs_tm :: Int
cs_tm = Int
p1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
p2, cs_ty :: Int
cs_ty = Int
q1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
q2, cs_co :: Int
cs_co = Int
r1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r2, cs_vb :: Int
cs_vb = Int
v1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
v2
       , cs_jb :: Int
cs_jb = Int
j1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
j2 }

zeroCS, oneTM :: CoreStats
zeroCS :: CoreStats
zeroCS = $WCS :: Int -> Int -> Int -> Int -> Int -> CoreStats
CS { cs_tm :: Int
cs_tm = 0, cs_ty :: Int
cs_ty = 0, cs_co :: Int
cs_co = 0, cs_vb :: Int
cs_vb = 0, cs_jb :: Int
cs_jb = 0 }
oneTM :: CoreStats
oneTM  = CoreStats
zeroCS { cs_tm :: Int
cs_tm = 1 }

sumCS :: (a -> CoreStats) -> [a] -> CoreStats
sumCS :: (a -> CoreStats) -> [a] -> CoreStats
sumCS f :: a -> CoreStats
f = (CoreStats -> a -> CoreStats) -> CoreStats -> [a] -> CoreStats
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\s :: CoreStats
s a :: a
a -> CoreStats -> CoreStats -> CoreStats
plusCS CoreStats
s (a -> CoreStats
f a
a)) CoreStats
zeroCS

coreBindsStats :: [CoreBind] -> CoreStats
coreBindsStats :: [CoreBind] -> CoreStats
coreBindsStats = (CoreBind -> CoreStats) -> [CoreBind] -> CoreStats
forall a. (a -> CoreStats) -> [a] -> CoreStats
sumCS (TopLevelFlag -> CoreBind -> CoreStats
bindStats TopLevelFlag
TopLevel)

bindStats :: TopLevelFlag -> CoreBind -> CoreStats
bindStats :: TopLevelFlag -> CoreBind -> CoreStats
bindStats top_lvl :: TopLevelFlag
top_lvl (NonRec v :: CoreBndr
v r :: Expr CoreBndr
r) = TopLevelFlag -> CoreBndr -> Expr CoreBndr -> CoreStats
bindingStats TopLevelFlag
top_lvl CoreBndr
v Expr CoreBndr
r
bindStats top_lvl :: TopLevelFlag
top_lvl (Rec prs :: [(CoreBndr, Expr CoreBndr)]
prs)    = ((CoreBndr, Expr CoreBndr) -> CoreStats)
-> [(CoreBndr, Expr CoreBndr)] -> CoreStats
forall a. (a -> CoreStats) -> [a] -> CoreStats
sumCS (\(v :: CoreBndr
v,r :: Expr CoreBndr
r) -> TopLevelFlag -> CoreBndr -> Expr CoreBndr -> CoreStats
bindingStats TopLevelFlag
top_lvl CoreBndr
v Expr CoreBndr
r) [(CoreBndr, Expr CoreBndr)]
prs

bindingStats :: TopLevelFlag -> Var -> CoreExpr -> CoreStats
bindingStats :: TopLevelFlag -> CoreBndr -> Expr CoreBndr -> CoreStats
bindingStats top_lvl :: TopLevelFlag
top_lvl v :: CoreBndr
v r :: Expr CoreBndr
r = TopLevelFlag -> CoreBndr -> CoreStats
letBndrStats TopLevelFlag
top_lvl CoreBndr
v CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
r

bndrStats :: Var -> CoreStats
bndrStats :: CoreBndr -> CoreStats
bndrStats v :: CoreBndr
v = CoreStats
oneTM CoreStats -> CoreStats -> CoreStats
`plusCS` Type -> CoreStats
tyStats (CoreBndr -> Type
varType CoreBndr
v)

letBndrStats :: TopLevelFlag -> Var -> CoreStats
letBndrStats :: TopLevelFlag -> CoreBndr -> CoreStats
letBndrStats top_lvl :: TopLevelFlag
top_lvl v :: CoreBndr
v
  | CoreBndr -> Bool
isTyVar CoreBndr
v Bool -> Bool -> Bool
|| TopLevelFlag -> Bool
isTopLevel TopLevelFlag
top_lvl = CoreBndr -> CoreStats
bndrStats CoreBndr
v
  | CoreBndr -> Bool
isJoinId CoreBndr
v = CoreStats
oneTM { cs_jb :: Int
cs_jb = 1 } CoreStats -> CoreStats -> CoreStats
`plusCS` CoreStats
ty_stats
  | Bool
otherwise  = CoreStats
oneTM { cs_vb :: Int
cs_vb = 1 } CoreStats -> CoreStats -> CoreStats
`plusCS` CoreStats
ty_stats
  where
    ty_stats :: CoreStats
ty_stats = Type -> CoreStats
tyStats (CoreBndr -> Type
varType CoreBndr
v)

exprStats :: CoreExpr -> CoreStats
exprStats :: Expr CoreBndr -> CoreStats
exprStats (Var {})        = CoreStats
oneTM
exprStats (Lit {})        = CoreStats
oneTM
exprStats (Type t :: Type
t)        = Type -> CoreStats
tyStats Type
t
exprStats (Coercion c :: Coercion
c)    = Coercion -> CoreStats
coStats Coercion
c
exprStats (App f :: Expr CoreBndr
f a :: Expr CoreBndr
a)       = Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
f CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
a
exprStats (Lam b :: CoreBndr
b e :: Expr CoreBndr
e)       = CoreBndr -> CoreStats
bndrStats CoreBndr
b CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
e
exprStats (Let b :: CoreBind
b e :: Expr CoreBndr
e)       = TopLevelFlag -> CoreBind -> CoreStats
bindStats TopLevelFlag
NotTopLevel CoreBind
b CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
e
exprStats (Case e :: Expr CoreBndr
e b :: CoreBndr
b _ as :: [Alt CoreBndr]
as) = Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
e CoreStats -> CoreStats -> CoreStats
`plusCS` CoreBndr -> CoreStats
bndrStats CoreBndr
b
                                        CoreStats -> CoreStats -> CoreStats
`plusCS` (Alt CoreBndr -> CoreStats) -> [Alt CoreBndr] -> CoreStats
forall a. (a -> CoreStats) -> [a] -> CoreStats
sumCS Alt CoreBndr -> CoreStats
altStats [Alt CoreBndr]
as
exprStats (Cast e :: Expr CoreBndr
e co :: Coercion
co)     = Coercion -> CoreStats
coStats Coercion
co CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
e
exprStats (Tick _ e :: Expr CoreBndr
e)      = Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
e

altStats :: CoreAlt -> CoreStats
altStats :: Alt CoreBndr -> CoreStats
altStats (_, bs :: [CoreBndr]
bs, r :: Expr CoreBndr
r) = [CoreBndr] -> CoreStats
altBndrStats [CoreBndr]
bs CoreStats -> CoreStats -> CoreStats
`plusCS` Expr CoreBndr -> CoreStats
exprStats Expr CoreBndr
r

altBndrStats :: [Var] -> CoreStats
-- Charge one for the alternative, not for each binder
altBndrStats :: [CoreBndr] -> CoreStats
altBndrStats vs :: [CoreBndr]
vs = CoreStats
oneTM CoreStats -> CoreStats -> CoreStats
`plusCS` (CoreBndr -> CoreStats) -> [CoreBndr] -> CoreStats
forall a. (a -> CoreStats) -> [a] -> CoreStats
sumCS (Type -> CoreStats
tyStats (Type -> CoreStats) -> (CoreBndr -> Type) -> CoreBndr -> CoreStats
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBndr -> Type
varType) [CoreBndr]
vs

tyStats :: Type -> CoreStats
tyStats :: Type -> CoreStats
tyStats ty :: Type
ty = CoreStats
zeroCS { cs_ty :: Int
cs_ty = Type -> Int
typeSize Type
ty }

coStats :: Coercion -> CoreStats
coStats :: Coercion -> CoreStats
coStats co :: Coercion
co = CoreStats
zeroCS { cs_co :: Int
cs_co = Coercion -> Int
coercionSize Coercion
co }

coreBindsSize :: [CoreBind] -> Int
-- We use coreBindStats for user printout
-- but this one is a quick and dirty basis for
-- the simplifier's tick limit
coreBindsSize :: [CoreBind] -> Int
coreBindsSize bs :: [CoreBind]
bs = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((CoreBind -> Int) -> [CoreBind] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map CoreBind -> Int
bindSize [CoreBind]
bs)

exprSize :: CoreExpr -> Int
-- ^ A measure of the size of the expressions, strictly greater than 0
-- Counts *leaves*, not internal nodes. Types and coercions are not counted.
exprSize :: Expr CoreBndr -> Int
exprSize (Var _)         = 1
exprSize (Lit _)         = 1
exprSize (App f :: Expr CoreBndr
f a :: Expr CoreBndr
a)       = Expr CoreBndr -> Int
exprSize Expr CoreBndr
f Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
a
exprSize (Lam b :: CoreBndr
b e :: Expr CoreBndr
e)       = CoreBndr -> Int
bndrSize CoreBndr
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e
exprSize (Let b :: CoreBind
b e :: Expr CoreBndr
e)       = CoreBind -> Int
bindSize CoreBind
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e
exprSize (Case e :: Expr CoreBndr
e b :: CoreBndr
b _ as :: [Alt CoreBndr]
as) = Expr CoreBndr -> Int
exprSize Expr CoreBndr
e Int -> Int -> Int
forall a. Num a => a -> a -> a
+ CoreBndr -> Int
bndrSize CoreBndr
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Alt CoreBndr -> Int) -> [Alt CoreBndr] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Alt CoreBndr -> Int
altSize [Alt CoreBndr]
as)
exprSize (Cast e :: Expr CoreBndr
e _)      = 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e
exprSize (Tick n :: Tickish CoreBndr
n e :: Expr CoreBndr
e)      = Tickish CoreBndr -> Int
tickSize Tickish CoreBndr
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e
exprSize (Type _)        = 1
exprSize (Coercion _)    = 1

tickSize :: Tickish Id -> Int
tickSize :: Tickish CoreBndr -> Int
tickSize (ProfNote _ _ _) = 1
tickSize _ = 1

bndrSize :: Var -> Int
bndrSize :: CoreBndr -> Int
bndrSize _ = 1

bndrsSize :: [Var] -> Int
bndrsSize :: [CoreBndr] -> Int
bndrsSize = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([CoreBndr] -> [Int]) -> [CoreBndr] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CoreBndr -> Int) -> [CoreBndr] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map CoreBndr -> Int
bndrSize

bindSize :: CoreBind -> Int
bindSize :: CoreBind -> Int
bindSize (NonRec b :: CoreBndr
b e :: Expr CoreBndr
e) = CoreBndr -> Int
bndrSize CoreBndr
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e
bindSize (Rec prs :: [(CoreBndr, Expr CoreBndr)]
prs)    = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((CoreBndr, Expr CoreBndr) -> Int)
-> [(CoreBndr, Expr CoreBndr)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, Expr CoreBndr) -> Int
pairSize [(CoreBndr, Expr CoreBndr)]
prs)

pairSize :: (Var, CoreExpr) -> Int
pairSize :: (CoreBndr, Expr CoreBndr) -> Int
pairSize (b :: CoreBndr
b,e :: Expr CoreBndr
e) = CoreBndr -> Int
bndrSize CoreBndr
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e

altSize :: CoreAlt -> Int
altSize :: Alt CoreBndr -> Int
altSize (_,bs :: [CoreBndr]
bs,e :: Expr CoreBndr
e) = [CoreBndr] -> Int
bndrsSize [CoreBndr]
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Expr CoreBndr -> Int
exprSize Expr CoreBndr
e