OverloadedCalls.hs

{-# LANGUAGE LambdaCase    #-}
{-# LANGUAGE TupleSections #-}

module GHC.Core.LateCC.OverloadedCalls
  ( overloadedCallsCC
  ) where

import GHC.Prelude

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

import GHC.Data.FastString
import GHC.Core
import GHC.Core.LateCC.Utils
import GHC.Core.LateCC.Types
import GHC.Core.Make
import GHC.Core.Predicate
import GHC.Core.Type
import GHC.Core.Utils
import GHC.Types.Id
import GHC.Types.Name
import GHC.Types.SrcLoc
import GHC.Types.Tickish
import GHC.Types.Var

type OverloadedCallsCCState = Strict.Maybe SrcSpan

{- Note [Overloaded Calls and join points]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Currently GHC considers cost centres as destructive to
join contexts. Or in other words this is not considered valid:

    join f x = ...
    in
              ... -> scc<tick> jmp

This makes the functionality of `-fprof-late-overloaded-calls` not feasible
for join points in general. We used to try to work around this by putting the
ticks on the rhs of the join point rather than around the jump. However beyond
the loss of accuracy this was broken for recursive join points as we ended up
with something like:

    rec-join f x = scc<tick> ... jmp f x

Which similarly is not valid as the tick once again destroys the tail call.
One might think we could limit ourselves to non-recursive tail calls and do
something clever like:

  join f x = scc<tick> ...
  in ... jmp f x

And sometimes this works! But sometimes the full rhs would look something like:

  join g x = ....
  join f x = scc<tick> ... -> jmp g x

Which, would again no longer be valid. I believe in the long run we can make
cost centre ticks non-destructive to join points. Or we could keep track of
where we are/are not allowed to insert a cost centre. But in the short term I will
simply disable the annotation of join calls under this flag.
-}

-- | Insert cost centres on function applications with dictionary arguments. The
-- source locations attached to the cost centres is approximated based on the
-- "closest" source note encountered in the traversal.
overloadedCallsCC :: CoreBind -> LateCCM OverloadedCallsCCState CoreBind
overloadedCallsCC =
    processBind
  where
    processBind :: CoreBind -> LateCCM OverloadedCallsCCState CoreBind
    processBind core_bind =
        case core_bind of
          NonRec b e ->
            NonRec b <$> wrap_if_join b (processExpr e)
          Rec es ->
            Rec <$> mapM (\(b,e) -> (b,) <$> wrap_if_join b (processExpr e)) es
      where
        -- If an overloaded function is turned into a join point, we won't add
        -- SCCs directly to calls since it makes them non-tail calls. Instead,
        -- we look for join points here and add an SCC to their RHS if they are
        -- overloaded.
        wrap_if_join ::
             CoreBndr
          -> LateCCM OverloadedCallsCCState CoreExpr
          -> LateCCM OverloadedCallsCCState CoreExpr
        wrap_if_join _b pexpr = do
            -- See Note [Overloaded Calls and join points]
            expr <- pexpr
            return expr

    processExpr :: CoreExpr -> LateCCM OverloadedCallsCCState CoreExpr
    processExpr expr =
      case expr of
        -- The case we care about: Application
        app@App{} -> do
          -- Here we have some application like `f v1 ... vN`, where v1 ... vN
          -- should be the function's type arguments followed by the value
          -- arguments. To determine if the `f` is an overloaded function, we
          -- check if any of the arguments v1 ... vN are dictionaries.
          let
            (f, xs) = collectArgs app
            resultTy = applyTypeToArgs (exprType f) xs

          -- Recursively process the arguments first for no particular reason
          args <- mapM processExpr xs
          let app' = mkCoreApps f args

          if
              -- Check if any of the arguments are dictionaries
              any isDictExpr args

              -- Avoid instrumenting dictionary functions, which may be
              -- overloaded if there are superclasses, by checking if the result
              -- type of the function is a dictionary type.
            && not (isDictTy resultTy)

              -- Avoid instrumenting constraint selectors like eq_sel
            && (typeTypeOrConstraint resultTy /= ConstraintLike)

              -- Avoid instrumenting join points.
              -- (See comment in processBind above)
              -- Also see Note [Overloaded Calls and join points]
            && not (isJoinVarExpr f)
          then do
            -- Extract a name and source location from the function being
            -- applied
            let
              cc_name :: FastString
              cc_name =
                maybe (fsLit "<no name available>") getOccFS (exprName app)

            cc_srcspan <-
              fmap (Strict.fromMaybe (UnhelpfulSpan UnhelpfulNoLocationInfo)) $
                lift $ gets lateCCState_extra

            insertCC cc_name cc_srcspan app'
          else
            return app'

        -- For recursive constructors of Expr, we traverse the nested Exprs
        Lam b e ->
          mkCoreLams [b] <$> processExpr e
        Let b e ->
          mkCoreLet <$> processBind b <*> processExpr e
        Case e b t alts ->
              Case
          <$> processExpr e
          <*> pure b
          <*> pure t
          <*> mapM processAlt alts
        Cast e co ->
          mkCast <$> processExpr e <*> pure co
        Tick t e -> do
          trackSourceNote t $
            mkTick t <$> processExpr e

        -- For non-recursive constructors of Expr, we do nothing
        x -> return x

    processAlt :: CoreAlt -> LateCCM OverloadedCallsCCState CoreAlt
    processAlt (Alt c bs e) = Alt c bs <$> processExpr e

    trackSourceNote :: CoreTickish -> LateCCM OverloadedCallsCCState a -> LateCCM OverloadedCallsCCState a
    trackSourceNote tick act =
      case tick of
        SourceNote rss _ -> do
          -- Prefer source notes from the current file
          in_current_file <-
            maybe False ((== EQ) . lexicalCompareFS (srcSpanFile rss)) <$>
              asks lateCCEnv_file
          if not in_current_file then
            act
          else do
            loc <- lift $ gets lateCCState_extra
            lift . modify $ \s ->
              s { lateCCState_extra =
                    Strict.Just $ RealSrcSpan rss mempty
                }
            x <- act
            lift . modify $ \s ->
              s { lateCCState_extra = loc
                }
            return x
        _ ->
          act

    -- Utility functions

    -- Extract a Name from an expression. If it is an application, attempt to
    -- extract a name from the applied function. If it is a variable, return the
    -- Name of the variable. If it is a tick/cast, attempt to extract a Name
    -- from the expression held in the tick/cast. Otherwise return Nothing.
    exprName :: CoreExpr -> Maybe Name
    exprName =
        \case
          App f _ ->
            exprName f
          Var f ->
            Just (idName f)
          Tick _ e ->
            exprName e
          Cast e _ ->
            exprName e
          _ ->
            Nothing

    -- Determine whether an expression is a dictionary
    isDictExpr :: CoreExpr -> Bool
    isDictExpr =
        maybe False isDictTy . exprType'
      where
        exprType' :: CoreExpr -> Maybe Type
        exprType' = \case
            Type{} -> Nothing
            expr -> Just $ exprType expr

    -- Determine whether an expression is a join variable
    isJoinVarExpr :: CoreExpr -> Bool
    isJoinVarExpr =
        \case
          Var var -> isJoinId var
          Tick _ e -> isJoinVarExpr e
          Cast e _ -> isJoinVarExpr e
          _ -> False