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

    {- *********************************************************************
    *                                                                      *
                             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 }