{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE TupleSections #-}

-- | Adds cost-centers to call sites selected with the @-fprof-caller=...@
-- flag.
module GHC.Core.Opt.CallerCC
    ( addCallerCostCentres
    , CallerCcFilter(..)
    , NamePattern(..)
    , parseCallerCcFilter
    ) where

import Data.Maybe

import Control.Applicative
import GHC.Utils.Monad.State.Strict
import Control.Monad

import GHC.Prelude
import GHC.Utils.Outputable as Outputable
import GHC.Driver.DynFlags
import GHC.Types.CostCentre
import GHC.Types.CostCentre.State
import GHC.Types.Name hiding (varName)
import GHC.Types.Tickish
import GHC.Unit.Module.ModGuts
import GHC.Types.SrcLoc
import GHC.Types.Var
import GHC.Unit.Types
import GHC.Data.FastString
import GHC.Core
import GHC.Core.Opt.Monad
import GHC.Core.Opt.CallerCC.Types


addCallerCostCentres :: ModGuts -> CoreM ModGuts
addCallerCostCentres :: ModGuts -> CoreM ModGuts
addCallerCostCentres ModGuts
guts = do
  dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  let filters = DynFlags -> [CallerCcFilter]
callerCcFilters DynFlags
dflags
  let env :: Env
      env = Env
        { thisModule :: Module
thisModule = ModGuts -> Module
mg_module ModGuts
guts
        , ccState :: CostCentreState
ccState = CostCentreState
newCostCentreState
        , countEntries :: Bool
countEntries = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_ProfCountEntries DynFlags
dflags
        , revParents :: [Id]
revParents = []
        , filters :: [CallerCcFilter]
filters = [CallerCcFilter]
filters
        }
  let guts' = ModGuts
guts { mg_binds = doCoreProgram env (mg_binds guts)
                   }
  return guts'

doCoreProgram :: Env -> CoreProgram -> CoreProgram
doCoreProgram :: Env -> CoreProgram -> CoreProgram
doCoreProgram Env
env CoreProgram
binds = (State CostCentreState CoreProgram
 -> CostCentreState -> CoreProgram)
-> CostCentreState
-> State CostCentreState CoreProgram
-> CoreProgram
forall a b c. (a -> b -> c) -> b -> a -> c
flip State CostCentreState CoreProgram -> CostCentreState -> CoreProgram
forall s a. State s a -> s -> a
evalState CostCentreState
newCostCentreState (State CostCentreState CoreProgram -> CoreProgram)
-> State CostCentreState CoreProgram -> CoreProgram
forall a b. (a -> b) -> a -> b
$ do
    (CoreBind -> State CostCentreState CoreBind)
-> CoreProgram -> State CostCentreState CoreProgram
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> CoreBind -> State CostCentreState CoreBind
doBind Env
env) CoreProgram
binds

doBind :: Env -> CoreBind -> M CoreBind
doBind :: Env -> CoreBind -> State CostCentreState CoreBind
doBind Env
env (NonRec Id
b Expr Id
rhs) = Id -> Expr Id -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
b (Expr Id -> CoreBind)
-> State CostCentreState (Expr Id)
-> State CostCentreState CoreBind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr (Id -> Env -> Env
addParent Id
b Env
env) Expr Id
rhs
doBind Env
env (Rec [(Id, Expr Id)]
bs) = [(Id, Expr Id)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([(Id, Expr Id)] -> CoreBind)
-> State CostCentreState [(Id, Expr Id)]
-> State CostCentreState CoreBind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id, Expr Id) -> State CostCentreState (Id, Expr Id))
-> [(Id, Expr Id)] -> State CostCentreState [(Id, Expr Id)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Id, Expr Id) -> State CostCentreState (Id, Expr Id)
doPair [(Id, Expr Id)]
bs
  where
    doPair :: (Id, Expr Id) -> State CostCentreState (Id, Expr Id)
doPair (Id
b,Expr Id
rhs) = (Id
b,) (Expr Id -> (Id, Expr Id))
-> State CostCentreState (Expr Id)
-> State CostCentreState (Id, Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr (Id -> Env -> Env
addParent Id
b Env
env) Expr Id
rhs

doExpr :: Env -> CoreExpr -> M CoreExpr
doExpr :: Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env e :: Expr Id
e@(Var Id
v)
  | Env -> Id -> Bool
needsCallSiteCostCentre Env
env Id
v = do
    let nameDoc :: SDoc
        nameDoc :: SDoc
nameDoc = NamePprCtx -> Depth -> SDoc -> SDoc
withUserStyle NamePprCtx
alwaysQualify Depth
DefaultDepth (SDoc -> SDoc) -> SDoc -> SDoc
forall a b. (a -> b) -> a -> b
$
          [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat (SDoc -> [SDoc] -> [SDoc]
forall doc. IsLine doc => doc -> [doc] -> [doc]
punctuate SDoc
forall doc. IsLine doc => doc
dot ((Id -> SDoc) -> [Id] -> [SDoc]
forall a b. (a -> b) -> [a] -> [b]
map Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr (Env -> [Id]
parents Env
env))) SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc
parens (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"calling:" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr Id
v)

        ccName :: CcName
        ccName :: CcName
ccName = String -> CcName
mkFastString (String -> CcName) -> String -> CcName
forall a b. (a -> b) -> a -> b
$ SDocContext -> SDoc -> String
renderWithContext SDocContext
defaultSDocContext SDoc
nameDoc
    ccIdx <- CcName -> M CostCentreIndex
getCCIndex' CcName
ccName
    let count = Env -> Bool
countEntries Env
env
        span = case Env -> [Id]
revParents Env
env of
          Id
top:[Id]
_ -> Name -> SrcSpan
nameSrcSpan (Name -> SrcSpan) -> Name -> SrcSpan
forall a b. (a -> b) -> a -> b
$ Id -> Name
varName Id
top
          [Id]
_     -> SrcSpan
noSrcSpan
        cc = CCFlavour -> CcName -> Module -> SrcSpan -> CostCentre
NormalCC (CostCentreIndex -> CCFlavour
mkExprCCFlavour CostCentreIndex
ccIdx) CcName
ccName (Env -> Module
thisModule Env
env) SrcSpan
span
        tick :: CoreTickish
        tick = CostCentre -> Bool -> Bool -> CoreTickish
forall (pass :: TickishPass).
CostCentre -> Bool -> Bool -> GenTickish pass
ProfNote CostCentre
cc Bool
count Bool
True
    pure $ Tick tick e
  | Bool
otherwise = Expr Id -> State CostCentreState (Expr Id)
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr Id
e
doExpr Env
_env e :: Expr Id
e@(Lit Literal
_)       = Expr Id -> State CostCentreState (Expr Id)
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr Id
e
doExpr Env
env (Expr Id
f `App` Expr Id
x)      = Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Expr Id -> Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
f State CostCentreState (Expr Id -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
x
doExpr Env
env (Lam Id
b Expr Id
x)        = Id -> Expr Id -> Expr Id
forall b. b -> Expr b -> Expr b
Lam Id
b (Expr Id -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
x
doExpr Env
env (Let CoreBind
b Expr Id
rhs)      = CoreBind -> Expr Id -> Expr Id
forall b. Bind b -> Expr b -> Expr b
Let (CoreBind -> Expr Id -> Expr Id)
-> State CostCentreState CoreBind
-> State CostCentreState (Expr Id -> Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> CoreBind -> State CostCentreState CoreBind
doBind Env
env CoreBind
b State CostCentreState (Expr Id -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
rhs
doExpr Env
env (Case Expr Id
scrut Id
b Type
ty [Alt Id]
alts) =
    Expr Id -> Id -> Type -> [Alt Id] -> Expr Id
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Expr Id -> Id -> Type -> [Alt Id] -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Id -> Type -> [Alt Id] -> Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
scrut State CostCentreState (Id -> Type -> [Alt Id] -> Expr Id)
-> State CostCentreState Id
-> State CostCentreState (Type -> [Alt Id] -> Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Id -> State CostCentreState Id
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Id
b State CostCentreState (Type -> [Alt Id] -> Expr Id)
-> State CostCentreState Type
-> State CostCentreState ([Alt Id] -> Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> State CostCentreState Type
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty State CostCentreState ([Alt Id] -> Expr Id)
-> State CostCentreState [Alt Id]
-> State CostCentreState (Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Alt Id -> State CostCentreState (Alt Id))
-> [Alt Id] -> State CostCentreState [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Alt Id -> State CostCentreState (Alt Id)
doAlt [Alt Id]
alts
  where
    doAlt :: Alt Id -> State CostCentreState (Alt Id)
doAlt (Alt AltCon
con [Id]
bs Expr Id
rhs)  = AltCon -> [Id] -> Expr Id -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con [Id]
bs (Expr Id -> Alt Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Alt Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
rhs
doExpr Env
env (Cast Expr Id
expr CoercionR
co)   = Expr Id -> CoercionR -> Expr Id
forall b. Expr b -> CoercionR -> Expr b
Cast (Expr Id -> CoercionR -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (CoercionR -> Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
expr State CostCentreState (CoercionR -> Expr Id)
-> State CostCentreState CoercionR
-> State CostCentreState (Expr Id)
forall a b.
State CostCentreState (a -> b)
-> State CostCentreState a -> State CostCentreState b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CoercionR -> State CostCentreState CoercionR
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CoercionR
co
doExpr Env
env (Tick CoreTickish
t Expr Id
e)       = CoreTickish -> Expr Id -> Expr Id
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t (Expr Id -> Expr Id)
-> State CostCentreState (Expr Id)
-> State CostCentreState (Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Expr Id -> State CostCentreState (Expr Id)
doExpr Env
env Expr Id
e
doExpr Env
_env e :: Expr Id
e@(Type Type
_)      = Expr Id -> State CostCentreState (Expr Id)
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr Id
e
doExpr Env
_env e :: Expr Id
e@(Coercion CoercionR
_)  = Expr Id -> State CostCentreState (Expr Id)
forall a. a -> State CostCentreState a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr Id
e

type M = State CostCentreState

getCCIndex' :: FastString -> M CostCentreIndex
getCCIndex' :: CcName -> M CostCentreIndex
getCCIndex' CcName
name = (CostCentreState -> (CostCentreIndex, CostCentreState))
-> M CostCentreIndex
forall s a. (s -> (a, s)) -> State s a
state (CcName -> CostCentreState -> (CostCentreIndex, CostCentreState)
getCCIndex CcName
name)

data Env = Env
  { Env -> Module
thisModule  :: Module
  , Env -> Bool
countEntries :: !Bool
  , Env -> CostCentreState
ccState     :: CostCentreState
  , Env -> [Id]
revParents  :: [Id]
  , Env -> [CallerCcFilter]
filters     :: [CallerCcFilter]
  }

addParent :: Id -> Env -> Env
addParent :: Id -> Env -> Env
addParent Id
i Env
env = Env
env { revParents = i : revParents env }

parents :: Env -> [Id]
parents :: Env -> [Id]
parents Env
env = [Id] -> [Id]
forall a. [a] -> [a]
reverse (Env -> [Id]
revParents Env
env)

needsCallSiteCostCentre :: Env -> Id -> Bool
needsCallSiteCostCentre :: Env -> Id -> Bool
needsCallSiteCostCentre Env
env Id
i =
    (CallerCcFilter -> Bool) -> [CallerCcFilter] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any CallerCcFilter -> Bool
matches (Env -> [CallerCcFilter]
filters Env
env)
  where
    matches :: CallerCcFilter -> Bool
    matches :: CallerCcFilter -> Bool
matches CallerCcFilter
ccf =
        Bool
checkModule Bool -> Bool -> Bool
&& Bool
checkFunc
      where
        checkModule :: Bool
checkModule =
          case CallerCcFilter -> Maybe ModuleName
ccfModuleName CallerCcFilter
ccf of
            Just ModuleName
modFilt
              | Just Module
iMod <- Name -> Maybe Module
nameModule_maybe (Id -> Name
varName Id
i)
              -> Module -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName Module
iMod ModuleName -> ModuleName -> Bool
forall a. Eq a => a -> a -> Bool
== ModuleName
modFilt
              | Bool
otherwise -> Bool
False
            Maybe ModuleName
Nothing -> Bool
True
        checkFunc :: Bool
checkFunc =
            NamePattern -> OccName -> Bool
occNameMatches (CallerCcFilter -> NamePattern
ccfFuncName CallerCcFilter
ccf) (Id -> OccName
forall a. NamedThing a => a -> OccName
getOccName Id
i)