{-# LANGUAGE TypeFamilies, DataKinds, GADTs, FlexibleInstances #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}

{-# LANGUAGE UndecidableInstances #-}
 -- To permit: type instance XLet 'InferTaggedBinders = XLet 'CodeGen

module GHC.Stg.InferTags.Types
    ( module GHC.Stg.InferTags.Types
    , module TagSig)
where

import GHC.Prelude

import GHC.Core.DataCon
import GHC.Core.Type (isUnliftedType, definitelyUnliftedType)
import GHC.Types.Id
import GHC.Stg.Syntax
import GHC.Stg.InferTags.TagSig as TagSig
import GHC.Types.Var.Env
import GHC.Utils.Outputable
import GHC.Utils.Misc( zipWithEqual )
import GHC.Utils.Panic

import GHC.StgToCmm.Types
import GHC.StgToCmm.Closure (importedIdLFInfo)

{- *********************************************************************
*                                                                      *
                         Supporting data types
*                                                                      *
********************************************************************* -}

type InferStgTopBinding = GenStgTopBinding 'InferTaggedBinders
type InferStgBinding    = GenStgBinding    'InferTaggedBinders
type InferStgExpr       = GenStgExpr       'InferTaggedBinders
type InferStgRhs        = GenStgRhs        'InferTaggedBinders
type InferStgAlt        = GenStgAlt        'InferTaggedBinders

-- | Combine TagInfo from the alternatives of a case expression.
-- Note that this operates at the value level: case alternatives return
-- values. See Note [TagSig and TagInfo].
combineAltInfo :: TagInfo -> TagInfo -> TagInfo
combineAltInfo :: TagInfo -> TagInfo -> TagInfo
combineAltInfo TagInfo
TagBottoming    TagInfo
ti              = TagInfo
ti
combineAltInfo TagInfo
ti              TagInfo
TagBottoming    = TagInfo
ti
combineAltInfo TagInfo
TagDunno        TagInfo
TagDunno        = TagInfo
TagDunno
combineAltInfo TagInfo
TagDunno        TagInfo
TagEPT          = TagInfo
TagDunno
combineAltInfo TagInfo
TagDunno        (TagTuple {})   = TagInfo
TagDunno
combineAltInfo TagInfo
TagEPT          TagInfo
TagDunno        = TagInfo
TagDunno
combineAltInfo (TagTuple {})   TagInfo
TagDunno        = TagInfo
TagDunno
combineAltInfo TagInfo
TagEPT          TagInfo
TagEPT          = TagInfo
TagEPT
-- TagEPT/TagTuple are incompatible (can arise with rep-polymorphic
-- results, see #26107); fall through to TagDunno.
combineAltInfo TagInfo
TagEPT          (TagTuple {})   = TagInfo
TagDunno
combineAltInfo (TagTuple {})   TagInfo
TagEPT          = TagInfo
TagDunno
combineAltInfo (TagTuple [TagInfo]
is1)  (TagTuple [TagInfo]
is2)  = [TagInfo] -> TagInfo
TagTuple (String
-> (TagInfo -> TagInfo -> TagInfo)
-> [TagInfo]
-> [TagInfo]
-> [TagInfo]
forall a b c.
HasDebugCallStack =>
String -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithEqual String
"combineAltInfo:TagTuple" TagInfo -> TagInfo -> TagInfo
combineAltInfo [TagInfo]
is1 [TagInfo]
is2)

type TagSigEnv = IdEnv TagSig
data TagEnv p = TE { forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env :: TagSigEnv
                   , forall (p :: StgPass). TagEnv p -> BinderP p -> Id
te_get :: BinderP p -> Id
                   , forall (p :: StgPass). TagEnv p -> Bool
te_bytecode :: !Bool
                   }

instance Outputable (TagEnv p) where
    ppr :: TagEnv p -> SDoc
ppr TagEnv p
te = SDoc
for_txt SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> TagSigEnv -> SDoc
forall a. Outputable a => a -> SDoc
ppr (TagEnv p -> TagSigEnv
forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env TagEnv p
te)
        where
            for_txt :: SDoc
for_txt = if TagEnv p -> Bool
forall (p :: StgPass). TagEnv p -> Bool
te_bytecode TagEnv p
te
                then String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"for_bytecode"
                else String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"for_native"

getBinderId :: TagEnv p -> BinderP p -> Id
getBinderId :: forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId = TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
te_get

initEnv :: Bool -> TagEnv 'CodeGen
initEnv :: Bool -> TagEnv 'CodeGen
initEnv Bool
for_bytecode = TE { te_env :: TagSigEnv
te_env = TagSigEnv
forall a. VarEnv a
emptyVarEnv
             , te_get :: BinderP 'CodeGen -> Id
te_get = \BinderP 'CodeGen
x -> Id
BinderP 'CodeGen
x
             , te_bytecode :: Bool
te_bytecode = Bool
for_bytecode }

-- | Simple convert env to a env of the 'InferTaggedBinders pass
-- with no other changes.
makeTagged :: TagEnv p -> TagEnv 'InferTaggedBinders
makeTagged :: forall (p :: StgPass). TagEnv p -> TagEnv 'InferTaggedBinders
makeTagged TagEnv p
env = TE { te_env :: TagSigEnv
te_env = TagEnv p -> TagSigEnv
forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env TagEnv p
env
                    , te_get :: BinderP 'InferTaggedBinders -> Id
te_get = (Id, TagSig) -> Id
BinderP 'InferTaggedBinders -> Id
forall a b. (a, b) -> a
fst
                    , te_bytecode :: Bool
te_bytecode = TagEnv p -> Bool
forall (p :: StgPass). TagEnv p -> Bool
te_bytecode TagEnv p
env }

noSig :: TagEnv p -> BinderP p -> (Id, TagSig)
noSig :: forall (p :: StgPass). TagEnv p -> BinderP p -> (Id, TagSig)
noSig TagEnv p
env BinderP p
bndr
  | HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType (Id -> Type
idType Id
var) = (Id
var, TagInfo -> TagSig
TagVal TagInfo
TagEPT)
  | Bool
otherwise = (Id
var, TagInfo -> TagSig
TagVal TagInfo
TagDunno)
  where
    var :: Id
var = TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
env BinderP p
bndr

-- | Look up the return-value tag of a function for saturated call analysis.
-- Returns 'Just retInfo' if the function binding has a 'TagFun' signature,
-- 'Nothing' otherwise.
lookupReturnInfo :: TagEnv p -> Id -> Maybe TagInfo
lookupReturnInfo :: forall (p :: StgPass). TagEnv p -> Id -> Maybe TagInfo
lookupReturnInfo TagEnv p
env Id
fun = case TagSigEnv -> Id -> Maybe TagSig
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (TagEnv p -> TagSigEnv
forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env TagEnv p
env) Id
fun of
  Just (TagFun TagInfo
ret_info) -> TagInfo -> Maybe TagInfo
forall a. a -> Maybe a
Just TagInfo
ret_info
  Just (TagVal TagInfo
_)        -> Maybe TagInfo
forall a. Maybe a
Nothing
  Maybe TagSig
Nothing                -> Maybe TagInfo
forall a. Maybe a
Nothing

-- | Look up a value-level tag.
--
-- If the tag info is completely determined by Id info we use that tag.
-- If the tag info is dependent on the analysis/call site use the passed TagSig.
-- If the passed TagSig is Nothing we fall back to safe tag info based on id properties.
argTagInfo :: Bool -> (Maybe TagSig) -> StgArg -> TagInfo
{-# INLINE argTagInfo #-} -- Specialize for the call site.
argTagInfo :: Bool -> Maybe TagSig -> StgArg -> TagInfo
argTagInfo Bool
_for_bytecode Maybe TagSig
_env_sig (StgLitArg{}) = TagInfo
TagEPT
argTagInfo Bool
for_bytecode Maybe TagSig
env_sig (StgVarArg Id
var)
  -- Nullary data constructors like True, False
  | Just DataCon
dc <- Id -> Maybe DataCon
isDataConWorkId_maybe Id
var
  , DataCon -> Bool
isNullaryRepDataCon DataCon
dc
  , Bool -> Bool
not Bool
for_bytecode
  = TagInfo
TagEPT

  -- NB: v might be the Id of a representation-polymorphic join point,
  -- so we shouldn't use isUnliftedType here. See T22212.
  | Type -> Bool
definitelyUnliftedType (Id -> Type
idType Id
var)
  = TagInfo
TagEPT

  -- Variables in the environment. A function binding flattens to TagEPT
  -- since a function closure pointer is properly tagged; its return info
  -- is not relevant when the function is used as a value.
  | Just TagSig
sig <- Maybe TagSig
env_sig
  = case TagSig
sig of
      TagVal TagInfo
info -> TagInfo
info
      TagFun TagInfo
_    -> TagInfo
TagEPT

-- Always save defaults.
  | LambdaFormInfo
lf_info <- Id -> LambdaFormInfo
importedIdLFInfo Id
var -- will make up LFInfo if not present.
  = if Bool
for_bytecode
    then TagInfo
TagDunno -- See Note [EPT enforcement for interpreted code] in GHC.Stg.EnforceEpt, point b)
    else
        case LambdaFormInfo
lf_info of
          -- Function, tagged (with arity)
          LFReEntrant {}
              -> TagInfo
TagEPT
          -- Thunks need to be entered.
          LFThunk {}
              -> TagInfo
TagDunno
          -- Constructors, already tagged.
          LFCon {}
              -> TagInfo
TagEPT
          LFUnknown {}
              -> TagInfo
TagDunno
          LFUnlifted {}
              -> TagInfo
TagEPT
          -- Shouldn't be possible. I don't think we can export letNoEscapes
          LFLetNoEscape {} -> String -> TagInfo
forall a. HasCallStack => String -> a
panic String
"LFLetNoEscape exported"

-- | Get tag info using the environment.
lookupInfo :: TagEnv p -> StgArg -> TagInfo
lookupInfo :: forall (p :: StgPass). TagEnv p -> StgArg -> TagInfo
lookupInfo TagEnv p
env arg :: StgArg
arg@(StgVarArg Id
var) =
    let env_sig :: Maybe TagSig
env_sig = (TagSigEnv -> Id -> Maybe TagSig
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (TagEnv p -> TagSigEnv
forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env TagEnv p
env) Id
var)
    in  Bool -> Maybe TagSig -> StgArg -> TagInfo
argTagInfo (TagEnv p -> Bool
forall (p :: StgPass). TagEnv p -> Bool
te_bytecode TagEnv p
env) Maybe TagSig
env_sig StgArg
arg
lookupInfo TagEnv p
_env (StgLitArg{}) = TagInfo
TagEPT

isDunnoSig :: TagSig -> Bool
isDunnoSig :: TagSig -> Bool
isDunnoSig (TagVal TagInfo
TagDunno) = Bool
True
isDunnoSig TagSig
_                 = Bool
False

-- | Extend the tag environment.
extendSigEnv :: TagEnv p -> [(Id,TagSig)] -> TagEnv p
extendSigEnv :: forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv env :: TagEnv p
env@(TE { te_env :: forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env = TagSigEnv
sig_env }) [(Id, TagSig)]
bndrs
  = TagEnv p
env { te_env = extendVarEnvList sig_env bndrs }