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

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

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

import GHC.Prelude

import GHC.Core.DataCon
import GHC.Core.Type (isUnliftedType)
import GHC.Types.Id
import GHC.Stg.Syntax
import GHC.Stg.EnforceEpt.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

{- *********************************************************************
*                                                                      *
                         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 ((TagInfo -> TagInfo -> TagInfo)
-> [TagInfo] -> [TagInfo] -> [TagInfo]
forall a b c.
HasDebugCallStack =>
(a -> b -> c) -> [a] -> [b] -> [c]
zipWithEqual 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 for an argument: either from the env (where
-- a function-typed argument flattens to TagEPT — its closure is tagged)
-- or derived from information on the variable itself.
lookupInfo :: TagEnv p -> StgArg -> TagInfo
lookupInfo :: forall (p :: StgPass). TagEnv p -> StgArg -> TagInfo
lookupInfo TagEnv p
env (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

  | HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType (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 <- 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
  = case TagSig
sig of
      TagVal TagInfo
info -> TagInfo
info
      TagFun TagInfo
_    -> TagInfo
TagEPT

  | Just LambdaFormInfo
lf_info <- Id -> Maybe LambdaFormInfo
idLFInfo_maybe Id
var
  , Bool -> Bool
not Bool
for_bytecode
  =   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"

  | Bool
otherwise
  = TagInfo
TagDunno
  where
    for_bytecode :: Bool
for_bytecode = TagEnv p -> Bool
forall (p :: StgPass). TagEnv p -> Bool
te_bytecode TagEnv p
env

lookupInfo TagEnv p
_ (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 }