    {-# 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)