1 files changed, 223 insertions, 0 deletions
diff --git a/compiler/GHC/Core/Opt/CallerCC.hs b/compiler/GHC/Core/Opt/CallerCC.hs
new file mode 100644
index 0000000000..1bbf96ca73
--- /dev/null
+++ b/compiler/GHC/Core/Opt/CallerCC.hs
@@ -0,0 +1,223 @@
+{-# 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
+    , parseCallerCcFilter
+    ) where
+
+import Data.Bifunctor
+import Data.Word (Word8)
+import Data.Maybe
+import qualified Text.Parsec as P
+
+import Control.Applicative
+import Control.Monad.Trans.State.Strict
+import Data.Either
+import Control.Monad
+
+import GHC.Prelude
+import GHC.Utils.Outputable as Outputable
+import GHC.Driver.Session
+import GHC.Driver.Ppr
+import GHC.Types.CostCentre
+import GHC.Types.CostCentre.State
+import GHC.Types.Name hiding (varName)
+import GHC.Unit.Module.Name
+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.Utils.Panic
+import qualified GHC.Utils.Binary as B
+
+addCallerCostCentres :: ModGuts -> CoreM ModGuts
+addCallerCostCentres guts = do
+  dflags <- getDynFlags
+  let filters = callerCcFilters dflags
+  let env :: Env
+      env = Env
+        { thisModule = mg_module guts
+        , ccState = newCostCentreState
+        , dflags = dflags
+        , revParents = []
+        , filters = filters
+        }
+  let guts' = guts { mg_binds = doCoreProgram env (mg_binds guts)
+                   }
+  return guts'
+
+doCoreProgram :: Env -> CoreProgram -> CoreProgram
+doCoreProgram env binds = flip evalState newCostCentreState $ do
+    mapM (doBind env) binds
+
+doBind :: Env -> CoreBind -> M CoreBind
+doBind env (NonRec b rhs) = NonRec b <$> doExpr (addParent b env) rhs
+doBind env (Rec bs) = Rec <$> mapM doPair bs
+  where
+    doPair (b,rhs) = (b,) <$> doExpr (addParent b env) rhs
+
+doExpr :: Env -> CoreExpr -> M CoreExpr
+doExpr env e@(Var v)
+  | needsCallSiteCostCentre env v = do
+    let nameDoc :: SDoc
+        nameDoc = withUserStyle alwaysQualify DefaultDepth $
+          hcat (punctuate dot (map ppr (parents env))) <> parens (text "calling:" <> ppr v)
+
+        ccName :: CcName
+        ccName = mkFastString $ showSDoc (dflags env) nameDoc
+    ccIdx <- getCCIndex' ccName
+    let span = case revParents env of
+          top:_ -> nameSrcSpan $ varName top
+          _     -> noSrcSpan
+        cc = NormalCC (ExprCC ccIdx) ccName (thisModule env) span
+        tick :: Tickish Id
+        tick = ProfNote cc True True
+    pure $ Tick tick e
+  | otherwise = pure e
+doExpr _env e@(Lit _)       = pure e
+doExpr env (f `App` x)      = App <$> doExpr env f <*> doExpr env x
+doExpr env (Lam b x)        = Lam b <$> doExpr env x
+doExpr env (Let b rhs)      = Let <$> doBind env b <*> doExpr env rhs
+doExpr env (Case scrut b ty alts) =
+    Case <$> doExpr env scrut <*> pure b <*> pure ty <*> mapM doAlt alts
+  where
+    doAlt (con, bs, rhs) = (con, bs,) <$> doExpr env rhs
+doExpr env (Cast expr co)   = Cast <$> doExpr env expr <*> pure co
+doExpr env (Tick t e)       = Tick t <$> doExpr env e
+doExpr _env e@(Type _)      = pure e
+doExpr _env e@(Coercion _)  = pure e
+
+type M = State CostCentreState
+
+getCCIndex' :: FastString -> M CostCentreIndex
+getCCIndex' name = state (getCCIndex name)
+
+data Env = Env
+  { thisModule  :: Module
+  , dflags      :: DynFlags
+  , ccState     :: CostCentreState
+  , revParents  :: [Id]
+  , filters     :: [CallerCcFilter]
+  }
+
+addParent :: Id -> Env -> Env
+addParent i env = env { revParents = i : revParents env }
+
+parents :: Env -> [Id]
+parents env = reverse (revParents env)
+
+needsCallSiteCostCentre :: Env -> Id -> Bool
+needsCallSiteCostCentre env i =
+    any matches (filters env)
+  where
+    matches :: CallerCcFilter -> Bool
+    matches ccf =
+        checkModule && checkFunc
+      where
+        checkModule =
+          case ccfModuleName ccf of
+            Just modFilt
+              | Just iMod <- nameModule_maybe (varName i)
+              -> moduleName iMod == modFilt
+              | otherwise -> False
+            Nothing -> True
+        checkFunc =
+            occNameMatches (ccfFuncName ccf) (getOccName i)
+
+data NamePattern
+    = PChar Char NamePattern
+    | PWildcard NamePattern
+    | PEnd
+
+instance Outputable NamePattern where
+  ppr (PChar c rest) = char c <> ppr rest
+  ppr (PWildcard rest) = char '*' <> ppr rest
+  ppr PEnd = Outputable.empty
+
+instance B.Binary NamePattern where
+  get bh = do
+    tag <- B.get bh
+    case tag :: Word8 of
+      0 -> PChar <$> B.get bh <*> B.get bh
+      1 -> PWildcard <$> B.get bh
+      2 -> pure PEnd
+      _ -> panic "Binary(NamePattern): Invalid tag"
+  put_ bh (PChar x y) = B.put_ bh (0 :: Word8) >> B.put_ bh x >> B.put_ bh y
+  put_ bh (PWildcard x) = B.put_ bh (1 :: Word8) >> B.put_ bh x
+  put_ bh PEnd = B.put_ bh (2 :: Word8)
+
+occNameMatches :: NamePattern -> OccName -> Bool
+occNameMatches pat = go pat . occNameString
+  where
+    go :: NamePattern -> String -> Bool
+    go PEnd "" = True
+    go (PChar c rest) (d:s)
+      = d == c && go rest s
+    go (PWildcard rest) s
+      = go rest s || go (PWildcard rest) (tail s)
+    go _ _  = False
+
+type Parser = P.Parsec String ()
+
+parseNamePattern :: Parser NamePattern
+parseNamePattern = pattern
+  where
+    pattern = star <|> wildcard <|> char <|> end
+    star = PChar '*' <$ P.string "\\*" <*> pattern
+    wildcard = do
+      void $ P.char '*'
+      PWildcard <$> pattern
+    char = PChar <$> P.anyChar <*> pattern
+    end = PEnd <$ P.eof
+
+data CallerCcFilter
+    = CallerCcFilter { ccfModuleName  :: Maybe ModuleName
+                     , ccfFuncName    :: NamePattern
+                     }
+
+instance Outputable CallerCcFilter where
+  ppr ccf =
+    maybe (char '*') ppr (ccfModuleName ccf)
+    <> char '.'
+    <> ppr (ccfFuncName ccf)
+
+instance B.Binary CallerCcFilter where
+  get bh = CallerCcFilter <$> B.get bh <*> B.get bh
+  put_ bh (CallerCcFilter x y) = B.put_ bh x >> B.put_ bh y
+
+parseCallerCcFilter :: String -> Either String CallerCcFilter
+parseCallerCcFilter =
+    first show . P.parse parseCallerCcFilter' "caller-CC filter"
+
+parseCallerCcFilter' :: Parser CallerCcFilter
+parseCallerCcFilter' =
+  CallerCcFilter
+    <$> moduleFilter
+    <*  P.char '.'
+    <*> parseNamePattern
+  where
+    moduleFilter :: Parser (Maybe ModuleName)
+    moduleFilter =
+      (Just . mkModuleName <$> moduleName)
+      <|>
+      (Nothing <$ P.char '*')
+
+    moduleName :: Parser String
+    moduleName = do
+      c <- P.upper
+      cs <- some $ P.upper <|> P.lower <|> P.digit <|> P.oneOf "_"
+      rest <- optional $ P.try $ P.char '.' >> fmap ('.':) moduleName
+      return $ c : (cs ++ fromMaybe "" rest)
+