{-# 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

combineAltInfo :: TagInfo -> TagInfo -> TagInfo
combineAltInfo :: TagInfo -> TagInfo -> TagInfo
combineAltInfo TagInfo
TagDunno         TagInfo
_              = TagInfo
TagDunno
combineAltInfo TagInfo
_                TagInfo
TagDunno       = TagInfo
TagDunno
combineAltInfo (TagTuple {})    TagInfo
TagProper      = String -> TagInfo
forall a. HasCallStack => String -> a
panic String
"Combining unboxed tuple with non-tuple result"
combineAltInfo TagInfo
TagProper       (TagTuple {})   = String -> TagInfo
forall a. HasCallStack => String -> a
panic String
"Combining unboxed tuple with non-tuple result"
combineAltInfo TagInfo
TagProper        TagInfo
TagProper      = TagInfo
TagProper
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" TagInfo -> TagInfo -> TagInfo
combineAltInfo [TagInfo]
is1 [TagInfo]
is2)
combineAltInfo (TagInfo
TagTagged)      TagInfo
ti             = TagInfo
ti
combineAltInfo TagInfo
ti               TagInfo
TagTagged      = TagInfo
ti

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
TagSig TagInfo
TagProper)
  | Bool
otherwise = (Id
var, TagInfo -> TagSig
TagSig 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 a sig in the given env
lookupSig :: TagEnv p -> Id -> Maybe TagSig
lookupSig :: forall (p :: StgPass). TagEnv p -> Id -> Maybe TagSig
lookupSig TagEnv p
env Id
fun = 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

-- | Look up a sig in the env or derive it from information
-- in the arg 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
TagProper

  | HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType (Id -> Type
idType Id
var)
  = TagInfo
TagProper

  -- Variables in the environment.
  | Just (TagSig TagInfo
info) <- 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
  = TagInfo
info

  | 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
TagProper
          -- Thunks need to be entered.
          LFThunk {}
              -> TagInfo
TagDunno
          -- Constructors, already tagged.
          LFCon {}
              -> TagInfo
TagProper
          LFUnknown {}
              -> TagInfo
TagDunno
          LFUnlifted {}
              -> TagInfo
TagProper
          -- 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
TagProper

isDunnoSig :: TagSig -> Bool
isDunnoSig :: TagSig -> Bool
isDunnoSig (TagSig TagInfo
TagDunno) = Bool
True
isDunnoSig (TagSig TagInfo
TagProper) = Bool
False
isDunnoSig (TagSig TagTuple{}) = Bool
False
isDunnoSig (TagSig TagTagged{}) = Bool
False

isTaggedInfo :: TagInfo -> Bool
isTaggedInfo :: TagInfo -> Bool
isTaggedInfo TagInfo
TagProper = Bool
True
isTaggedInfo TagInfo
TagTagged = Bool
True
isTaggedInfo TagInfo
_         = Bool
False

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 }