module GHC.Core.LateCC.Utils
  ( -- * Inserting cost centres
    doLateCostCenters -- Might be useful for API users

    -- ** Helpers for defining insertion methods
  , getCCFlavour
  , insertCC
  ) where

import GHC.Prelude

import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State.Strict
import qualified Data.Set as S

import GHC.Core
import GHC.Core.LateCC.Types
import GHC.Core.Utils
import GHC.Data.FastString
import GHC.Types.CostCentre
import GHC.Types.CostCentre.State
import GHC.Types.SrcLoc
import GHC.Types.Tickish

-- | Insert cost centres into the 'CoreProgram' using the provided environment,
-- initial state, and insertion method.
doLateCostCenters
  :: LateCCEnv
  -- ^ Environment to run the insertion in
  -> LateCCState s
  -- ^ Initial state to run the insertion with
  -> (CoreBind -> LateCCM s CoreBind)
  -- ^ Insertion method
  -> CoreProgram
  -- ^ Bindings to consider
  -> (CoreProgram, LateCCState s)
doLateCostCenters env state method binds =
    runLateCC env state $ mapM method binds

-- | Evaluate late cost centre insertion
runLateCC :: LateCCEnv -> LateCCState s -> LateCCM s a -> (a, LateCCState s)
runLateCC env state = (`runState` state) . (`runReaderT` env)

-- | Given the name of a cost centre, get its flavour
getCCFlavour :: FastString -> LateCCM s CCFlavour
getCCFlavour name = mkLateCCFlavour <$> getCCIndex' name
  where
    getCCIndex' :: FastString -> LateCCM s CostCentreIndex
    getCCIndex' name = do
      cc_state <- lift $ gets lateCCState_ccState
      let (index, cc_state') = getCCIndex name cc_state
      lift . modify $ \s -> s { lateCCState_ccState = cc_state'}
      return index

-- | Insert a cost centre with the specified name and source span on the given
-- expression. The inserted cost centre will be appropriately tracked in the
-- late cost centre state.
insertCC
  :: FastString
  -- ^ Name of the cost centre to insert
  -> SrcSpan
  -- ^ Source location to associate with the cost centre
  -> CoreExpr
  -- ^ Expression to wrap in the cost centre
  -> LateCCM s CoreExpr
insertCC cc_name cc_loc expr = do
    cc_flavour <- getCCFlavour cc_name
    env <- ask
    let
      cc_mod = lateCCEnv_module env
      cc = NormalCC cc_flavour cc_name cc_mod cc_loc
      note = ProfNote cc (lateCCEnv_countEntries env) True
    when (lateCCEnv_collectCCs env) $ do
        lift . modify $ \s ->
          s { lateCCState_ccs = S.insert cc (lateCCState_ccs s)
            }
    return $ mkTick note expr