{-# LANGUAGE TypeFamilies, DataKinds, GADTs, FlexibleInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
 -- To permit: type instance XLet 'InferTaggedBinders = XLet 'SomePass

{-# OPTIONS_GHC -Wname-shadowing #-}
module GHC.Stg.InferTags ( inferTags ) where

import GHC.Prelude hiding (id)

import GHC.Core.DataCon
import GHC.Core.Type
import GHC.Types.Id
import GHC.Types.Id.Info (tagSigInfo)
import GHC.Types.Name
import GHC.Stg.Syntax
import GHC.Types.Basic ( CbvMark (..) )
import GHC.Types.Unique.Supply (mkSplitUniqSupply)
import GHC.Types.RepType (dataConRuntimeRepStrictness)
import GHC.Core (AltCon(..))
import Data.List (mapAccumL)
import GHC.Utils.Outputable
import GHC.Utils.Misc( zipWithEqual, zipEqual, notNull )

import GHC.Stg.InferTags.Types
import GHC.Stg.InferTags.Rewrite (rewriteTopBinds)
import Data.Maybe
import GHC.Types.Name.Env (mkNameEnv, NameEnv)
import GHC.Driver.Config.Stg.Ppr
import GHC.Driver.Session
import GHC.Utils.Logger
import qualified GHC.Unit.Types

{- Note [Tag Inference]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The purpose of this pass is to attach to every binder a flag
to indicate whether or not it is "properly tagged".  A binder
is properly tagged if it is guaranteed:
 - to point to a heap-allocated *value*
 - and to have the tag of the value encoded in the pointer

For example
  let x = Just y in ...

Here x will be properly tagged: it will point to the heap-allocated
values for (Just y), and the tag-bits of the pointer will encode
the tag for Just so there is no need to re-enter the closure or even
check for the presence of tag bits. The impacts of this can be very large.

For containers the reduction in runtimes with this optimization was as follows:

intmap-benchmarks:    89.30%
intset-benchmarks:    90.87%
map-benchmarks:       88.00%
sequence-benchmarks:  99.84%
set-benchmarks:       85.00%
set-operations-intmap:88.64%
set-operations-map:   74.23%
set-operations-set:   76.50%
lookupge-intmap:      89.57%
lookupge-map:         70.95%

With nofib being ~0.3% faster as well.

See Note [Tag inference passes] for how we proceed to generate and use this information.

Note [Strict Field Invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
As part of tag inference we introduce the Strict Field Invariant.
Which consists of us saying that:

* Pointers in strict fields must (a) point directly to the value, and
  (b) must be properly tagged.

For example, given
  data T = MkT ![Int]

the Strict Field Invariant guarantees that the first field of any `MkT` constructor
will either point directly to nil, or directly to a cons cell;
and will be tagged with `001` or `010` respectively.
It will never point to a thunk, nor will it be tagged `000` (meaning "might be a thunk").
NB: Note that the proper tag for some objects is indeed `000`. Currently this is the case for PAPs.

This works analogous to how `StrictWorkerId`s work. See also Note [Strict Worker Ids].

Why do we care? Because if we have code like:

case strictPair of
  SP x y ->
    case x of ...

It allows us to safely omit the code to enter x and the check
for the presence of a tag that goes along with it.
However we might still branch on the tag as usual.
See Note [Tag Inference] for how much impact this can have for
some code.

This is enforced by the code GHC.Stg.InferTags.Rewrite
where we:

* Look at all constructor allocations.
* Check if arguments to their strict fields are known to be properly tagged
* If not we convert `StrictJust x` into `case x of x' -> StrictJust x'`

This is usually very beneficial but can cause regressions in rare edge cases where
we fail to proof that x is properly tagged, or where it simply isn't.
See Note [How untagged pointers can end up in strict fields] for how the second case
can arise.

For a full example of the worst case consider this code:

foo ... = ...
  let c = StrictJust x
  in ...

Here we would rewrite `let c = StrictJust x` into `let c = case x of x' -> StrictJust x'`
However that is horrible! We end up allocating a thunk for `c` first, which only when
evaluated will allocate the constructor.

So we do our best to establish that `x` is already tagged (which it almost always is)
to avoid this cost. In my benchmarks I haven't seen any cases where this causes regressions.

Note that there are similar constraints around Note [Strict Worker Ids].

Note [How untagged pointers can end up in strict fields]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  data Set a = Tip | Bin !a (Set a) (Set a)

We make a wrapper for Bin that evaluates its arguments
  $WBin x a b = case x of xv -> Bin xv a b
Here `xv` will always be evaluated and properly tagged, just as the
Strict Field Invariant requires.

But alas the Simplifier can destroy the invariant: see #15696.
We start with
  thk = f ()
  g x = ...(case thk of xv -> Bin xv Tip Tip)...

So far so good; the argument to Bin (which is strict) is evaluated.
Now we do float-out. And in doing so we do a reverse binder-swap (see
Note [Binder-swap during float-out] in SetLevels) thus

  g x = ...(case thk of xv -> Bin thk Nil Nil)...

The goal of the reverse binder-swap is to allow more floating -- and
indeed it does! We float the Bin to top level:

  lvl = Bin thk Tip Tip
  g x = ...(case thk of xv -> lvl)...

Now you can see that the argument of Bin, namely thk, points to the
thunk, not to the value as it did before.

In short, although it may be rare, the output of optimisation passes
cannot guarantee to obey the Strict Field Invariant. For this reason
we run tag inference. See Note [Tag inference passes].

Note [Tag inference passes]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Tag inference proceeds in two passes:
* The first pass is an analysis to compute which binders are properly tagged.
  The result is then attached to /binders/.
  This is implemented by `inferTagsAnal` in GHC.Stg.InferTags
* The second pass walks over the AST checking if the Strict Field Invariant is upheld.
  See Note [Strict Field Invariant].
  If required this pass modifies the program to uphold this invariant.
  Tag information is also moved from /binders/ to /occurrences/ during this pass.
  This is done by `GHC.Stg.InferTags.Rewrite (rewriteTopBinds)`.
* Finally the code generation uses this information to skip the thunk check when branching on
  values. This is done by `cgExpr`/`cgCase` in the backend.

Last but not least we also export the tag sigs of top level bindings to allow this optimization
 to work across module boundaries.

Note [TagInfo of functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
The purpose of tag inference is really to figure out when we don't have to enter
value closures. There the meaning of the tag is fairly obvious.
For functions we never make use of the tag info so we have two choices:
* Treat them as TagDunno
* Treat them as TagProper (as they *are* tagged with their arity) and be really
  careful to make sure we still enter them when needed.
As it makes little difference for runtime performance I've treated functions as TagDunno in a few places where
it made the code simpler. But besides implementation complexity there isn't any reason
why we couldn't be more rigourous in dealing with functions.

NB: It turned out because of #21193 option two wouldn't really have been an option anyway.

Note [Tag inference debugging]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
There is a flag -dtag-inference-checks which inserts various
compile/runtime checks in order to ensure the Strict Field Invariant
holds. It should cover all places
where tags matter and disable optimizations which interfere with checking
the invariant like generation of AP-Thunks.

Note [Polymorphic StgPass for inferTagExpr]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In order to reach a fixpoint we sometimes have to re-analyse an expression
multiple times. But after the initial run the Ast will be parameterized by
a different StgPass! To handle this a large part of the analysis is polymorphic
over the exact StgPass we are using. Which allows us to run the analysis on
the output of itself.

-}

{- *********************************************************************
*                                                                      *
                         Tag inference pass
*                                                                      *
********************************************************************* -}

-- doCodeGen :: HscEnv -> Module -> InfoTableProvMap -> [TyCon]
--           -> CollectedCCs
--           -> [CgStgTopBinding] -- ^ Bindings come already annotated with fvs
--           -> HpcInfo
--           -> IO (Stream IO CmmGroupSRTs CgInfos)
--          -- Note we produce a 'Stream' of CmmGroups, so that the
--          -- backend can be run incrementally.  Otherwise it generates all
--          -- the C-- up front, which has a significant space cost.
inferTags :: DynFlags -> Logger -> (GHC.Unit.Types.Module) -> [CgStgTopBinding] -> IO ([TgStgTopBinding], NameEnv TagSig)
inferTags :: DynFlags
-> Logger
-> Module
-> [CgStgTopBinding]
-> IO ([CgStgTopBinding], NameEnv TagSig)
inferTags DynFlags
dflags Logger
logger Module
this_mod [CgStgTopBinding]
stg_binds = do

    -- Annotate binders with tag information.
    let (![GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags) = {-# SCC "StgTagFields" #-}
                                        [CgStgTopBinding] -> [GenStgTopBinding 'InferTaggedBinders]
inferTagsAnal [CgStgTopBinding]
stg_binds
    Logger -> DumpFlag -> String -> DumpFormat -> SDoc -> IO ()
putDumpFileMaybe Logger
logger DumpFlag
Opt_D_dump_stg_tags String
"CodeGenAnal STG:" DumpFormat
FormatSTG (StgPprOpts -> [GenStgTopBinding 'InferTaggedBinders] -> SDoc
forall (pass :: StgPass).
OutputablePass pass =>
StgPprOpts -> [GenStgTopBinding pass] -> SDoc
pprGenStgTopBindings (DynFlags -> StgPprOpts
initStgPprOpts DynFlags
dflags) [GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags)

    let export_tag_info :: NameEnv TagSig
export_tag_info = [GenStgTopBinding 'InferTaggedBinders] -> NameEnv TagSig
collectExportInfo [GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags

    -- Rewrite STG to uphold the strict field invariant
    UniqSupply
us_t <- Char -> IO UniqSupply
mkSplitUniqSupply Char
't'
    let rewritten_binds :: [CgStgTopBinding]
rewritten_binds = {-# SCC "StgTagRewrite" #-} Module
-> UniqSupply
-> [GenStgTopBinding 'InferTaggedBinders]
-> [CgStgTopBinding]
rewriteTopBinds Module
this_mod UniqSupply
us_t [GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags :: [TgStgTopBinding]

    ([CgStgTopBinding], NameEnv TagSig)
-> IO ([CgStgTopBinding], NameEnv TagSig)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([CgStgTopBinding]
rewritten_binds,NameEnv TagSig
export_tag_info)

{- *********************************************************************
*                                                                      *
                         Main inference algorithm
*                                                                      *
********************************************************************* -}

type OutputableInferPass p = (Outputable (TagEnv p)
                              , Outputable (GenStgExpr p)
                              , Outputable (BinderP p)
                              , Outputable (GenStgRhs p))

-- | This constraint encodes the fact that no matter what pass
-- we use the Let/Closure extension points are the same as these for
-- 'InferTaggedBinders.
type InferExtEq i = ( XLet i ~ XLet 'InferTaggedBinders
                    , XLetNoEscape i ~ XLetNoEscape 'InferTaggedBinders
                    , XRhsClosure i ~ XRhsClosure 'InferTaggedBinders)

inferTagsAnal :: [GenStgTopBinding 'CodeGen] -> [GenStgTopBinding 'InferTaggedBinders]
inferTagsAnal :: [CgStgTopBinding] -> [GenStgTopBinding 'InferTaggedBinders]
inferTagsAnal [CgStgTopBinding]
binds =
  -- pprTrace "Binds" (pprGenStgTopBindings shortStgPprOpts $ binds) $
  (TagEnv 'CodeGen, [GenStgTopBinding 'InferTaggedBinders])
-> [GenStgTopBinding 'InferTaggedBinders]
forall a b. (a, b) -> b
snd ((TagEnv 'CodeGen
 -> CgStgTopBinding
 -> (TagEnv 'CodeGen, GenStgTopBinding 'InferTaggedBinders))
-> TagEnv 'CodeGen
-> [CgStgTopBinding]
-> (TagEnv 'CodeGen, [GenStgTopBinding 'InferTaggedBinders])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL TagEnv 'CodeGen
-> CgStgTopBinding
-> (TagEnv 'CodeGen, GenStgTopBinding 'InferTaggedBinders)
inferTagTopBind TagEnv 'CodeGen
initEnv [CgStgTopBinding]
binds)

-----------------------
inferTagTopBind :: TagEnv 'CodeGen -> GenStgTopBinding 'CodeGen
                -> (TagEnv 'CodeGen, GenStgTopBinding 'InferTaggedBinders)
inferTagTopBind :: TagEnv 'CodeGen
-> CgStgTopBinding
-> (TagEnv 'CodeGen, GenStgTopBinding 'InferTaggedBinders)
inferTagTopBind TagEnv 'CodeGen
env (StgTopStringLit Id
id ByteString
bs)
  = (TagEnv 'CodeGen
env, Id -> ByteString -> GenStgTopBinding 'InferTaggedBinders
forall (pass :: StgPass). Id -> ByteString -> GenStgTopBinding pass
StgTopStringLit Id
id ByteString
bs)
inferTagTopBind TagEnv 'CodeGen
env (StgTopLifted GenStgBinding 'CodeGen
bind)
  = (TagEnv 'CodeGen
env', GenStgBinding 'InferTaggedBinders
-> GenStgTopBinding 'InferTaggedBinders
forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted GenStgBinding 'InferTaggedBinders
bind')
  where
    (TagEnv 'CodeGen
env', GenStgBinding 'InferTaggedBinders
bind') = TagEnv 'CodeGen
-> GenStgBinding 'CodeGen
-> (TagEnv 'CodeGen, GenStgBinding 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
inferTagBind TagEnv 'CodeGen
env GenStgBinding 'CodeGen
bind


-- Why is this polymorphic over the StgPass? See Note [Polymorphic StgPass for inferTagExpr]
-----------------------
inferTagExpr :: forall p. (OutputableInferPass p, InferExtEq p)
  => TagEnv p -> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr :: forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
env (StgApp Id
fun [StgArg]
args)
  =  --pprTrace "inferTagExpr1"
      -- (ppr fun <+> ppr args $$ ppr info $$
      --  text "deadEndInfo:" <> ppr (isDeadEndId fun, idArity fun, length args)
      -- )
    (TagInfo
info, Id -> [StgArg] -> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
fun [StgArg]
args)
  where
    !fun_arity :: Arity
fun_arity = Id -> Arity
idArity Id
fun
    info :: TagInfo
info | Arity
fun_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
0 -- Unknown arity => Thunk or unknown call
         = TagInfo
TagDunno

         | Id -> Bool
isDeadEndId Id
fun
         , Arity
fun_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== [StgArg] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [StgArg]
args -- Implies we will simply call the function.
         = TagInfo
TagTagged -- See Note [Bottom functions are TagTagged]

         | Just (TagSig TagInfo
res_info) <- IdInfo -> Maybe TagSig
tagSigInfo ((() :: Constraint) => Id -> IdInfo
Id -> IdInfo
idInfo Id
fun)
         , Arity
fun_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== [StgArg] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [StgArg]
args  -- Saturated
         = TagInfo
res_info

         | Just (TagSig TagInfo
res_info) <- TagEnv p -> Id -> Maybe TagSig
forall (p :: StgPass). TagEnv p -> Id -> Maybe TagSig
lookupSig TagEnv p
env Id
fun
         , Arity
fun_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== [StgArg] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [StgArg]
args  -- Saturated
         = TagInfo
res_info

         | Bool
otherwise
         = --pprTrace "inferAppUnknown" (ppr fun) $
           TagInfo
TagDunno
-- TODO:
-- If we have something like:
--   let x = thunk in
--   f g = case g of g' -> (# x, g' #)
-- then we *do* know that g' will be properly tagged,
-- so we should return TagTagged [TagDunno,TagProper] but currently we infer
-- TagTagged [TagDunno,TagDunno] because of the unknown arity case in inferTagExpr.
-- Seems not to matter much but should be changed eventually.

inferTagExpr TagEnv p
env (StgConApp DataCon
con ConstructorNumber
cn [StgArg]
args [Type]
tys)
  = (TagEnv p -> DataCon -> [StgArg] -> TagInfo
forall (p :: StgPass). TagEnv p -> DataCon -> [StgArg] -> TagInfo
inferConTag TagEnv p
env DataCon
con [StgArg]
args, DataCon
-> ConstructorNumber
-> [StgArg]
-> [Type]
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
DataCon
-> ConstructorNumber -> [StgArg] -> [Type] -> GenStgExpr pass
StgConApp DataCon
con ConstructorNumber
cn [StgArg]
args [Type]
tys)

inferTagExpr TagEnv p
_ (StgLit Literal
l)
  = (TagInfo
TagTagged, Literal -> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass). Literal -> GenStgExpr pass
StgLit Literal
l)

inferTagExpr TagEnv p
env (StgTick StgTickish
tick GenStgExpr p
body)
  = (TagInfo
info, StgTickish
-> GenStgExpr 'InferTaggedBinders -> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
StgTickish -> GenStgExpr pass -> GenStgExpr pass
StgTick StgTickish
tick GenStgExpr 'InferTaggedBinders
body')
  where
    (TagInfo
info, GenStgExpr 'InferTaggedBinders
body') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
env GenStgExpr p
body

inferTagExpr TagEnv p
_ (StgOpApp StgOp
op [StgArg]
args Type
ty)
  = -- Do any primops guarantee to return a properly tagged value?
    -- I think not.  Ditto foreign calls.
    (TagInfo
TagDunno, StgOp -> [StgArg] -> Type -> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
StgOp -> [StgArg] -> Type -> GenStgExpr pass
StgOpApp StgOp
op [StgArg]
args Type
ty)

inferTagExpr TagEnv p
env (StgLet XLet p
ext GenStgBinding p
bind GenStgExpr p
body)
  = (TagInfo
info, XLet 'InferTaggedBinders
-> GenStgBinding 'InferTaggedBinders
-> GenStgExpr 'InferTaggedBinders
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet p
XLet 'InferTaggedBinders
ext GenStgBinding 'InferTaggedBinders
bind' GenStgExpr 'InferTaggedBinders
body')
  where
    (TagEnv p
env', GenStgBinding 'InferTaggedBinders
bind') = TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
inferTagBind TagEnv p
env GenStgBinding p
bind
    (TagInfo
info, GenStgExpr 'InferTaggedBinders
body') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
env' GenStgExpr p
body

inferTagExpr TagEnv p
env (StgLetNoEscape XLetNoEscape p
ext GenStgBinding p
bind GenStgExpr p
body)
  = (TagInfo
info, XLetNoEscape 'InferTaggedBinders
-> GenStgBinding 'InferTaggedBinders
-> GenStgExpr 'InferTaggedBinders
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape p
XLetNoEscape 'InferTaggedBinders
ext GenStgBinding 'InferTaggedBinders
bind' GenStgExpr 'InferTaggedBinders
body')
  where
    (TagEnv p
env', GenStgBinding 'InferTaggedBinders
bind') = TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
inferTagBind TagEnv p
env GenStgBinding p
bind
    (TagInfo
info, GenStgExpr 'InferTaggedBinders
body') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
env' GenStgExpr p
body

inferTagExpr TagEnv p
in_env (StgCase GenStgExpr p
scrut BinderP p
bndr AltType
ty [GenStgAlt p]
alts)
  -- Unboxed tuples get their info from the expression we scrutinise if any
  | [GenStgAlt{alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con=DataAlt DataCon
con, alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs=[BinderP p]
bndrs, alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs=GenStgExpr p
rhs}] <- [GenStgAlt p]
alts
  , DataCon -> Bool
isUnboxedTupleDataCon DataCon
con
  , Just [TagInfo]
infos <- [BinderP p] -> Maybe [TagInfo]
scrut_infos [BinderP p]
bndrs
  , let bndrs' :: [(Id, TagSig)]
bndrs' = String
-> (BinderP p -> TagInfo -> (Id, TagSig))
-> [BinderP p]
-> [TagInfo]
-> [(Id, TagSig)]
forall a b c. String -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithEqual String
"inferTagExpr" BinderP p -> TagInfo -> (Id, TagSig)
mk_bndr [BinderP p]
bndrs [TagInfo]
infos
        mk_bndr :: BinderP p -> TagInfo -> (Id, TagSig)
        mk_bndr :: BinderP p -> TagInfo -> (Id, TagSig)
mk_bndr BinderP p
tup_bndr TagInfo
tup_info =
            --  pprTrace "mk_ubx_bndr_info" ( ppr bndr <+> ppr info ) $
            (TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
in_env BinderP p
tup_bndr, TagInfo -> TagSig
TagSig TagInfo
tup_info)
        -- no case binder in alt_env here, unboxed tuple binders are dead after unarise
        alt_env :: TagEnv p
alt_env = TagEnv p -> [(Id, TagSig)] -> TagEnv p
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv p
in_env [(Id, TagSig)]
bndrs'
        (TagInfo
info, GenStgExpr 'InferTaggedBinders
rhs') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
alt_env GenStgExpr p
rhs
  =
    -- pprTrace "inferCase1" (
    --   text "scrut:" <> ppr scrut $$
    --   text "bndr:" <> ppr bndr $$
    --   text "infos" <> ppr infos $$
    --   text "out_bndrs" <> ppr bndrs') $
    (TagInfo
info, GenStgExpr 'InferTaggedBinders
-> BinderP 'InferTaggedBinders
-> AltType
-> [GenStgAlt 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'InferTaggedBinders
scrut' (TagEnv p -> BinderP p -> (Id, TagSig)
forall (p :: StgPass). TagEnv p -> BinderP p -> (Id, TagSig)
noSig TagEnv p
in_env BinderP p
bndr) AltType
ty [GenStgAlt{ alt_con :: AltCon
alt_con=DataCon -> AltCon
DataAlt DataCon
con
                                                           , alt_bndrs :: [BinderP 'InferTaggedBinders]
alt_bndrs=[(Id, TagSig)]
[BinderP 'InferTaggedBinders]
bndrs'
                                                           , alt_rhs :: GenStgExpr 'InferTaggedBinders
alt_rhs=GenStgExpr 'InferTaggedBinders
rhs'}])

  | [GenStgAlt p] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [GenStgAlt p]
alts -- Empty case, but I might just be paranoid.
  = -- pprTrace "inferCase2" empty $
    (TagInfo
TagDunno, GenStgExpr 'InferTaggedBinders
-> BinderP 'InferTaggedBinders
-> AltType
-> [GenStgAlt 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'InferTaggedBinders
scrut' (Id, TagSig)
BinderP 'InferTaggedBinders
bndr' AltType
ty [])
  -- More than one alternative OR non-TagTuple single alternative.
  | Bool
otherwise
  =
    let
        case_env :: TagEnv p
case_env = TagEnv p -> [(Id, TagSig)] -> TagEnv p
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv p
in_env [(Id, TagSig)
bndr']

        ([TagInfo]
infos, [GenStgAlt 'InferTaggedBinders]
alts')
          = [(TagInfo, GenStgAlt 'InferTaggedBinders)]
-> ([TagInfo], [GenStgAlt 'InferTaggedBinders])
forall a b. [(a, b)] -> ([a], [b])
unzip [ (TagInfo
info, GenStgAlt p
g {alt_bndrs :: [BinderP 'InferTaggedBinders]
alt_bndrs=[BinderP 'InferTaggedBinders]
bndrs', alt_rhs :: GenStgExpr 'InferTaggedBinders
alt_rhs=GenStgExpr 'InferTaggedBinders
rhs'})
                  | g :: GenStgAlt p
g@GenStgAlt{ alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con = AltCon
con
                               , alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs = [BinderP p]
bndrs
                               , alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs   = GenStgExpr p
rhs
                               } <- [GenStgAlt p]
alts
                  , let (TagEnv p
alt_env,[BinderP 'InferTaggedBinders]
bndrs') = TagEnv p
-> AltCon
-> [BinderP p]
-> (TagEnv p, [BinderP 'InferTaggedBinders])
forall (p :: StgPass).
TagEnv p
-> AltCon
-> [BinderP p]
-> (TagEnv p, [BinderP 'InferTaggedBinders])
addAltBndrInfo TagEnv p
case_env AltCon
con [BinderP p]
bndrs
                        (TagInfo
info, GenStgExpr 'InferTaggedBinders
rhs') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
alt_env GenStgExpr p
rhs
                  ]
        alt_info :: TagInfo
alt_info = (TagInfo -> TagInfo -> TagInfo) -> TagInfo -> [TagInfo] -> TagInfo
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TagInfo -> TagInfo -> TagInfo
combineAltInfo TagInfo
TagTagged [TagInfo]
infos
    in ( TagInfo
alt_info, GenStgExpr 'InferTaggedBinders
-> BinderP 'InferTaggedBinders
-> AltType
-> [GenStgAlt 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'InferTaggedBinders
scrut' (Id, TagSig)
BinderP 'InferTaggedBinders
bndr' AltType
ty [GenStgAlt 'InferTaggedBinders]
alts')
  where
    -- Single unboxed tuple alternative
    scrut_infos :: [BinderP p] -> Maybe [TagInfo]
scrut_infos [BinderP p]
bndrs = case TagInfo
scrut_info of
      TagInfo
TagTagged -> [TagInfo] -> Maybe [TagInfo]
forall a. a -> Maybe a
Just ([TagInfo] -> Maybe [TagInfo]) -> [TagInfo] -> Maybe [TagInfo]
forall a b. (a -> b) -> a -> b
$ Arity -> TagInfo -> [TagInfo]
forall a. Arity -> a -> [a]
replicate ([BinderP p] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [BinderP p]
bndrs) TagInfo
TagProper
      TagTuple [TagInfo]
infos -> [TagInfo] -> Maybe [TagInfo]
forall a. a -> Maybe a
Just [TagInfo]
infos
      TagInfo
_ -> Maybe [TagInfo]
forall a. Maybe a
Nothing
    (TagInfo
scrut_info, GenStgExpr 'InferTaggedBinders
scrut') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
in_env GenStgExpr p
scrut
    bndr' :: (Id, TagSig)
bndr' = (TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
in_env BinderP p
bndr, TagInfo -> TagSig
TagSig TagInfo
TagProper)

-- Compute binder sigs based on the constructors strict fields.
-- NB: Not used if we have tuple info from the scrutinee.
addAltBndrInfo :: forall p. TagEnv p -> AltCon -> [BinderP p] -> (TagEnv p, [BinderP 'InferTaggedBinders])
addAltBndrInfo :: forall (p :: StgPass).
TagEnv p
-> AltCon
-> [BinderP p]
-> (TagEnv p, [BinderP 'InferTaggedBinders])
addAltBndrInfo TagEnv p
env (DataAlt DataCon
con) [BinderP p]
bndrs
  | Bool -> Bool
not (DataCon -> Bool
isUnboxedTupleDataCon DataCon
con Bool -> Bool -> Bool
|| DataCon -> Bool
isUnboxedSumDataCon DataCon
con)
  = (TagEnv p
out_env, [(Id, TagSig)]
[BinderP 'InferTaggedBinders]
out_bndrs)
  where
    marks :: [StrictnessMark]
marks = (() :: Constraint) => DataCon -> [StrictnessMark]
DataCon -> [StrictnessMark]
dataConRuntimeRepStrictness DataCon
con :: [StrictnessMark]
    out_bndrs :: [(Id, TagSig)]
out_bndrs = (BinderP p -> StrictnessMark -> (Id, TagSig))
-> [BinderP p] -> [StrictnessMark] -> [(Id, TagSig)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith BinderP p -> StrictnessMark -> (Id, TagSig)
mk_bndr [BinderP p]
bndrs [StrictnessMark]
marks
    out_env :: TagEnv p
out_env = TagEnv p -> [(Id, TagSig)] -> TagEnv p
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv p
env [(Id, TagSig)]
out_bndrs

    mk_bndr :: (BinderP p -> StrictnessMark -> (Id, TagSig))
    mk_bndr :: BinderP p -> StrictnessMark -> (Id, TagSig)
mk_bndr BinderP p
bndr StrictnessMark
mark
      | (() :: Constraint) => Type -> Bool
Type -> Bool
isUnliftedType (Id -> Type
idType Id
id) Bool -> Bool -> Bool
|| StrictnessMark -> Bool
isMarkedStrict StrictnessMark
mark
      = (Id
id, TagInfo -> TagSig
TagSig TagInfo
TagProper)
      | Bool
otherwise
      = TagEnv p -> BinderP p -> (Id, TagSig)
forall (p :: StgPass). TagEnv p -> BinderP p -> (Id, TagSig)
noSig TagEnv p
env BinderP p
bndr
        where
          id :: Id
id = TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
env BinderP p
bndr

addAltBndrInfo TagEnv p
env AltCon
_ [BinderP p]
bndrs = (TagEnv p
env, (BinderP p -> (Id, TagSig)) -> [BinderP p] -> [(Id, TagSig)]
forall a b. (a -> b) -> [a] -> [b]
map (TagEnv p -> BinderP p -> (Id, TagSig)
forall (p :: StgPass). TagEnv p -> BinderP p -> (Id, TagSig)
noSig TagEnv p
env) [BinderP p]
bndrs)

-----------------------------
inferTagBind :: (OutputableInferPass p, InferExtEq p)
  => TagEnv p -> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
inferTagBind :: forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
inferTagBind TagEnv p
in_env (StgNonRec BinderP p
bndr GenStgRhs p
rhs)
  =
    -- pprTrace "inferBindNonRec" (
    --   ppr bndr $$
    --   ppr (isDeadEndId id) $$
    --   ppr sig)
    (TagEnv p
env', BinderP 'InferTaggedBinders
-> GenStgRhs 'InferTaggedBinders
-> GenStgBinding 'InferTaggedBinders
forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec (Id
id, TagSig
sig) GenStgRhs 'InferTaggedBinders
rhs')
  where
    id :: Id
id   = TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
in_env BinderP p
bndr
    env' :: TagEnv p
env' = TagEnv p -> [(Id, TagSig)] -> TagEnv p
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv p
in_env [(Id
id, TagSig
sig)]
    (TagSig
sig,GenStgRhs 'InferTaggedBinders
rhs') = Id
-> TagEnv p
-> GenStgRhs p
-> (TagSig, GenStgRhs 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
Id
-> TagEnv p
-> GenStgRhs p
-> (TagSig, GenStgRhs 'InferTaggedBinders)
inferTagRhs Id
id TagEnv p
in_env GenStgRhs p
rhs

inferTagBind TagEnv p
in_env (StgRec [(BinderP p, GenStgRhs p)]
pairs)
  = -- pprTrace "rec" (ppr (map fst pairs) $$ ppr (in_env { te_env = out_env }, StgRec pairs')) $
    (TagEnv p
in_env { te_env :: TagSigEnv
te_env = TagSigEnv
out_env }, [(BinderP 'InferTaggedBinders, GenStgRhs 'InferTaggedBinders)]
-> GenStgBinding 'InferTaggedBinders
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [((Id, TagSig), GenStgRhs 'InferTaggedBinders)]
[(BinderP 'InferTaggedBinders, GenStgRhs 'InferTaggedBinders)]
pairs')
  where
    ([BinderP p]
bndrs, [GenStgRhs p]
rhss)     = [(BinderP p, GenStgRhs p)] -> ([BinderP p], [GenStgRhs p])
forall a b. [(a, b)] -> ([a], [b])
unzip [(BinderP p, GenStgRhs p)]
pairs
    in_ids :: [Id]
in_ids            = (BinderP p -> Id) -> [BinderP p] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
in_env) [BinderP p]
bndrs
    init_sigs :: [TagSig]
init_sigs         = ((Id, GenStgRhs p) -> TagSig) -> [(Id, GenStgRhs p)] -> [TagSig]
forall a b. (a -> b) -> [a] -> [b]
map ((Id, GenStgRhs p) -> TagSig
forall (p :: StgPass). (Id, GenStgRhs p) -> TagSig
initSig) ([(Id, GenStgRhs p)] -> [TagSig])
-> [(Id, GenStgRhs p)] -> [TagSig]
forall a b. (a -> b) -> a -> b
$ [Id] -> [GenStgRhs p] -> [(Id, GenStgRhs p)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
in_ids [GenStgRhs p]
rhss
    (TagSigEnv
out_env, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)]
pairs') = TagEnv p
-> [TagSig]
-> [GenStgRhs p]
-> (TagSigEnv, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)])
forall (q :: StgPass).
(OutputableInferPass q, InferExtEq q) =>
TagEnv q
-> [TagSig]
-> [GenStgRhs q]
-> (TagSigEnv, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)])
go TagEnv p
in_env [TagSig]
init_sigs [GenStgRhs p]
rhss

    go :: forall q. (OutputableInferPass q , InferExtEq q) => TagEnv q -> [TagSig] -> [GenStgRhs q]
                 -> (TagSigEnv, [((Id,TagSig), GenStgRhs 'InferTaggedBinders)])
    go :: forall (q :: StgPass).
(OutputableInferPass q, InferExtEq q) =>
TagEnv q
-> [TagSig]
-> [GenStgRhs q]
-> (TagSigEnv, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)])
go TagEnv q
go_env [TagSig]
in_sigs [GenStgRhs q]
go_rhss
      --   | pprTrace "go" (ppr in_ids $$ ppr in_sigs $$ ppr out_sigs $$ ppr rhss') False
      --  = undefined
       | [TagSig]
in_sigs [TagSig] -> [TagSig] -> Bool
forall a. Eq a => a -> a -> Bool
== [TagSig]
out_sigs = (TagEnv q -> TagSigEnv
forall (p :: StgPass). TagEnv p -> TagSigEnv
te_env TagEnv q
rhs_env, [(Id, TagSig)]
out_bndrs [(Id, TagSig)]
-> [GenStgRhs 'InferTaggedBinders]
-> [((Id, TagSig), GenStgRhs 'InferTaggedBinders)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [GenStgRhs 'InferTaggedBinders]
rhss')
       | Bool
otherwise     = TagEnv 'InferTaggedBinders
-> [TagSig]
-> [GenStgRhs 'InferTaggedBinders]
-> (TagSigEnv, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)])
forall (q :: StgPass).
(OutputableInferPass q, InferExtEq q) =>
TagEnv q
-> [TagSig]
-> [GenStgRhs q]
-> (TagSigEnv, [((Id, TagSig), GenStgRhs 'InferTaggedBinders)])
go TagEnv 'InferTaggedBinders
env' [TagSig]
out_sigs [GenStgRhs 'InferTaggedBinders]
rhss'
       where
         out_bndrs :: [(Id, TagSig)]
out_bndrs = ((Id, TagSig) -> (Id, TagSig)) -> [(Id, TagSig)] -> [(Id, TagSig)]
forall a b. (a -> b) -> [a] -> [b]
map (Id, TagSig) -> (Id, TagSig)
updateBndr [(Id, TagSig)]
in_bndrs -- TODO: Keeps in_ids alive
         in_bndrs :: [(Id, TagSig)]
in_bndrs = [Id]
in_ids [Id] -> [TagSig] -> [(Id, TagSig)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [TagSig]
in_sigs
         rhs_env :: TagEnv q
rhs_env = TagEnv q -> [(Id, TagSig)] -> TagEnv q
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv q
go_env [(Id, TagSig)]
in_bndrs
         ([TagSig]
out_sigs, [GenStgRhs 'InferTaggedBinders]
rhss') = [(TagSig, GenStgRhs 'InferTaggedBinders)]
-> ([TagSig], [GenStgRhs 'InferTaggedBinders])
forall a b. [(a, b)] -> ([a], [b])
unzip (String
-> (Id -> GenStgRhs q -> (TagSig, GenStgRhs 'InferTaggedBinders))
-> [Id]
-> [GenStgRhs q]
-> [(TagSig, GenStgRhs 'InferTaggedBinders)]
forall a b c. String -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithEqual String
"inferTagBind" Id -> GenStgRhs q -> (TagSig, GenStgRhs 'InferTaggedBinders)
anaRhs [Id]
in_ids [GenStgRhs q]
go_rhss)
         env' :: TagEnv 'InferTaggedBinders
env' = TagEnv q -> TagEnv 'InferTaggedBinders
forall (p :: StgPass). TagEnv p -> TagEnv 'InferTaggedBinders
makeTagged TagEnv q
go_env

         anaRhs :: Id -> GenStgRhs q -> (TagSig, GenStgRhs 'InferTaggedBinders)
         anaRhs :: Id -> GenStgRhs q -> (TagSig, GenStgRhs 'InferTaggedBinders)
anaRhs Id
bnd GenStgRhs q
rhs = Id
-> TagEnv q
-> GenStgRhs q
-> (TagSig, GenStgRhs 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
Id
-> TagEnv p
-> GenStgRhs p
-> (TagSig, GenStgRhs 'InferTaggedBinders)
inferTagRhs Id
bnd TagEnv q
rhs_env GenStgRhs q
rhs

         updateBndr :: (Id,TagSig) -> (Id,TagSig)
         updateBndr :: (Id, TagSig) -> (Id, TagSig)
updateBndr (Id
v,TagSig
sig) = (Id -> TagSig -> Id
setIdTagSig Id
v TagSig
sig, TagSig
sig)

initSig :: forall p. (Id, GenStgRhs p) -> TagSig
-- Initial signature for the fixpoint loop
initSig :: forall (p :: StgPass). (Id, GenStgRhs p) -> TagSig
initSig (Id
_bndr, StgRhsCon {})               = TagInfo -> TagSig
TagSig TagInfo
TagTagged
initSig (Id
bndr, StgRhsClosure XRhsClosure p
_ CostCentreStack
_ UpdateFlag
_ [BinderP p]
_ GenStgExpr p
_) =
  TagSig -> Maybe TagSig -> TagSig
forall a. a -> Maybe a -> a
fromMaybe TagSig
defaultSig (Id -> Maybe TagSig
idTagSig_maybe Id
bndr)
  where defaultSig :: TagSig
defaultSig = (TagInfo -> TagSig
TagSig TagInfo
TagTagged)

{- Note [Bottom functions are TagTagged]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If we have a function with two branches with one
being bottom, and the other returning a tagged
unboxed tuple what is the result? We give it TagTagged!
To answer why consider this function:

foo :: Bool -> (# Bool, Bool #)
foo x = case x of
    True -> (# True,True #)
    False -> undefined

The true branch is obviously tagged. The other branch isn't.
We want to treat the *result* of foo as tagged as well so that
the combination of the branches also is tagged if all non-bottom
branches are tagged.
This is safe because the function is still always called/entered as long
as it's applied to arguments. Since the function will never return we can give
it safely any tag sig we like.
So we give it TagTagged, as it allows the combined tag sig of the case expression
to be the combination of all non-bottoming branches.

-}

-----------------------------
inferTagRhs :: forall p.
     (OutputableInferPass p, InferExtEq p)
  => Id -- ^ Id we are binding to.
  -> TagEnv p -- ^
  -> GenStgRhs p -- ^
  -> (TagSig, GenStgRhs 'InferTaggedBinders)
inferTagRhs :: forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
Id
-> TagEnv p
-> GenStgRhs p
-> (TagSig, GenStgRhs 'InferTaggedBinders)
inferTagRhs Id
bnd_id TagEnv p
in_env (StgRhsClosure XRhsClosure p
ext CostCentreStack
cc UpdateFlag
upd [BinderP p]
bndrs GenStgExpr p
body)
  | Id -> Bool
isDeadEndId Id
bnd_id Bool -> Bool -> Bool
&& (forall (f :: * -> *) a. Foldable f => f a -> Bool
notNull) [BinderP p]
bndrs
  -- See Note [Bottom functions are TagTagged]
  = (TagInfo -> TagSig
TagSig TagInfo
TagTagged, XRhsClosure 'InferTaggedBinders
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> GenStgRhs pass
StgRhsClosure XRhsClosure p
XRhsClosure 'InferTaggedBinders
ext CostCentreStack
cc UpdateFlag
upd [(Id, TagSig)]
[BinderP 'InferTaggedBinders]
out_bndrs GenStgExpr 'InferTaggedBinders
body')
  | Bool
otherwise
  = --pprTrace "inferTagRhsClosure" (ppr (_top, _grp_ids, env,info')) $
    (TagInfo -> TagSig
TagSig TagInfo
info', XRhsClosure 'InferTaggedBinders
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> GenStgRhs pass
StgRhsClosure XRhsClosure p
XRhsClosure 'InferTaggedBinders
ext CostCentreStack
cc UpdateFlag
upd [(Id, TagSig)]
[BinderP 'InferTaggedBinders]
out_bndrs GenStgExpr 'InferTaggedBinders
body')
  where
    out_bndrs :: [(Id, TagSig)]
out_bndrs
      | Just [CbvMark]
marks <- Id -> Maybe [CbvMark]
idCbvMarks_maybe Id
bnd_id
      -- Sometimes an we eta-expand foo with additional arguments after ww, and we also trim
      -- the list of marks to the last strict entry. So we can conservatively
      -- assume these are not strict
      = (BinderP p -> CbvMark -> (Id, TagSig))
-> [BinderP p] -> [CbvMark] -> [(Id, TagSig)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (BinderP p -> CbvMark -> (Id, TagSig)
mkArgSig) [BinderP p]
bndrs ([CbvMark]
marks [CbvMark] -> [CbvMark] -> [CbvMark]
forall a. [a] -> [a] -> [a]
++ CbvMark -> [CbvMark]
forall a. a -> [a]
repeat CbvMark
NotMarkedCbv)
      | Bool
otherwise = (BinderP p -> (Id, TagSig)) -> [BinderP p] -> [(Id, TagSig)]
forall a b. (a -> b) -> [a] -> [b]
map (TagEnv p -> BinderP p -> (Id, TagSig)
forall (p :: StgPass). TagEnv p -> BinderP p -> (Id, TagSig)
noSig TagEnv p
env') [BinderP p]
bndrs :: [(Id,TagSig)]

    env' :: TagEnv p
env' = TagEnv p -> [(Id, TagSig)] -> TagEnv p
forall (p :: StgPass). TagEnv p -> [(Id, TagSig)] -> TagEnv p
extendSigEnv TagEnv p
in_env [(Id, TagSig)]
out_bndrs
    (TagInfo
info, GenStgExpr 'InferTaggedBinders
body') = TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
forall (p :: StgPass).
(OutputableInferPass p, InferExtEq p) =>
TagEnv p
-> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
inferTagExpr TagEnv p
env' GenStgExpr p
body
    info' :: TagInfo
info'
      -- It's a thunk
      | [BinderP p] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [BinderP p]
bndrs
      = TagInfo
TagDunno
      -- TODO: We could preserve tuple fields for thunks
      -- as well. But likely not worth the complexity.

      | Bool
otherwise  = TagInfo
info

    mkArgSig :: BinderP p -> CbvMark -> (Id,TagSig)
    mkArgSig :: BinderP p -> CbvMark -> (Id, TagSig)
mkArgSig BinderP p
bndp CbvMark
mark =
      let id :: Id
id = TagEnv p -> BinderP p -> Id
forall (p :: StgPass). TagEnv p -> BinderP p -> Id
getBinderId TagEnv p
in_env BinderP p
bndp
          tag :: TagInfo
tag = case CbvMark
mark of
            CbvMark
MarkedCbv -> TagInfo
TagProper
            CbvMark
_
              | (() :: Constraint) => Type -> Bool
Type -> Bool
isUnliftedType (Id -> Type
idType Id
id) -> TagInfo
TagProper
              | Bool
otherwise -> TagInfo
TagDunno
      in (Id
id, TagInfo -> TagSig
TagSig TagInfo
tag)

inferTagRhs Id
_ TagEnv p
env _rhs :: GenStgRhs p
_rhs@(StgRhsCon CostCentreStack
cc DataCon
con ConstructorNumber
cn [StgTickish]
ticks [StgArg]
args)
-- Constructors, which have untagged arguments to strict fields
-- become thunks. We encode this by giving changing RhsCon nodes the info TagDunno
  = --pprTrace "inferTagRhsCon" (ppr grp_ids) $
    (TagInfo -> TagSig
TagSig (TagEnv p -> DataCon -> [StgArg] -> TagInfo
forall (p :: StgPass). TagEnv p -> DataCon -> [StgArg] -> TagInfo
inferConTag TagEnv p
env DataCon
con [StgArg]
args), CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs pass
StgRhsCon CostCentreStack
cc DataCon
con ConstructorNumber
cn [StgTickish]
ticks [StgArg]
args)

{- Note [Constructor TagSigs]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@inferConTag@ will infer the proper tag signature for a binding who's RHS is a constructor
or a StgConApp expression.
Usually these will simply be TagProper. But there are exceptions.
If any of the fields in the constructor are strict, but any argument to these
fields is not tagged then we will have to case on the argument before storing
in the constructor. Which means for let bindings the RHS turns into a thunk
which obviously is no longer properly tagged.
For example we might start with:

    let x<TagDunno> = f ...
    let c<TagProper> = StrictPair x True

But we know during the rewrite stage x will need to be evaluated in the RHS
of `c` so we will infer:

    let x<TagDunno> = f ...
    let c<TagDunno> = StrictPair x True

Which in the rewrite stage will then be rewritten into:

    let x<TagDunno> = f ...
    let c<TagDunno> = case x of x' -> StrictPair x' True

The other exception is unboxed tuples. These will get a TagTuple
signature with a list of TagInfo about their individual binders
as argument. As example:

    let c<TagProper> = True
    let x<TagDunno> = ...
    let f<?> z = case z of z'<TagProper> -> (# c, x #)

Here we will infer for f the Signature <TagTuple[TagProper,TagDunno]>.
This information will be used if we scrutinze a saturated application of
`f` in order to determine the taggedness of the result.
That is for `case f x of (# r1,r2 #) -> rhs` we can infer
r1<TagProper> and r2<TagDunno> which allows us to skip all tag checks on `r1`
in `rhs`.

Things get a bit more complicated with nesting:

    let closeFd<TagTuple[...]> = ...
    let f x = ...
        case x of
          _ -> Solo# closeFd

The "natural" signature for the Solo# branch in `f` would be <TagTuple[TagTuple[...]]>.
But we flatten this out to <TagTuple[TagDunno]> for the time being as it improves compile
time and there doesn't seem to huge benefit to doing differently.

  -}

-- See Note [Constructor TagSigs]
inferConTag :: TagEnv p -> DataCon -> [StgArg] -> TagInfo
inferConTag :: forall (p :: StgPass). TagEnv p -> DataCon -> [StgArg] -> TagInfo
inferConTag TagEnv p
env DataCon
con [StgArg]
args
  | DataCon -> Bool
isUnboxedTupleDataCon DataCon
con
  = [TagInfo] -> TagInfo
TagTuple ([TagInfo] -> TagInfo) -> [TagInfo] -> TagInfo
forall a b. (a -> b) -> a -> b
$ (StgArg -> TagInfo) -> [StgArg] -> [TagInfo]
forall a b. (a -> b) -> [a] -> [b]
map (TagInfo -> TagInfo
flatten_arg_tag (TagInfo -> TagInfo) -> (StgArg -> TagInfo) -> StgArg -> TagInfo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TagEnv p -> StgArg -> TagInfo
forall (p :: StgPass). TagEnv p -> StgArg -> TagInfo
lookupInfo TagEnv p
env) [StgArg]
args
  | Bool
otherwise =
    -- pprTrace "inferConTag"
    --   ( text "con:" <> ppr con $$
    --     text "args:" <> ppr args $$
    --     text "marks:" <> ppr (dataConRuntimeRepStrictness con) $$
    --     text "arg_info:" <> ppr (map (lookupInfo env) args) $$
    --     text "info:" <> ppr info) $
    TagInfo
info
  where
    info :: TagInfo
info = if ((StgArg, StrictnessMark) -> Bool)
-> [(StgArg, StrictnessMark)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (StgArg, StrictnessMark) -> Bool
arg_needs_eval [(StgArg, StrictnessMark)]
strictArgs then TagInfo
TagDunno else TagInfo
TagProper
    strictArgs :: [(StgArg, StrictnessMark)]
strictArgs = String
-> [StgArg] -> [StrictnessMark] -> [(StgArg, StrictnessMark)]
forall a b. String -> [a] -> [b] -> [(a, b)]
zipEqual String
"inferTagRhs" [StgArg]
args ((() :: Constraint) => DataCon -> [StrictnessMark]
DataCon -> [StrictnessMark]
dataConRuntimeRepStrictness DataCon
con) :: ([(StgArg, StrictnessMark)])
    arg_needs_eval :: (StgArg, StrictnessMark) -> Bool
arg_needs_eval (StgArg
arg,StrictnessMark
strict)
      -- lazy args
      | Bool -> Bool
not (StrictnessMark -> Bool
isMarkedStrict StrictnessMark
strict) = Bool
False
      | TagInfo
tag <- (TagEnv p -> StgArg -> TagInfo
forall (p :: StgPass). TagEnv p -> StgArg -> TagInfo
lookupInfo TagEnv p
env StgArg
arg)
      -- banged args need to be tagged, or require eval
      = Bool -> Bool
not (TagInfo -> Bool
isTaggedInfo TagInfo
tag)

    flatten_arg_tag :: TagInfo -> TagInfo
flatten_arg_tag (TagInfo
TagTagged) = TagInfo
TagProper
    flatten_arg_tag (TagInfo
TagProper ) = TagInfo
TagProper
    flatten_arg_tag (TagTuple [TagInfo]
_) = TagInfo
TagDunno -- See Note [Constructor TagSigs]
    flatten_arg_tag (TagInfo
TagDunno) = TagInfo
TagDunno


collectExportInfo :: [GenStgTopBinding 'InferTaggedBinders] -> NameEnv TagSig
collectExportInfo :: [GenStgTopBinding 'InferTaggedBinders] -> NameEnv TagSig
collectExportInfo [GenStgTopBinding 'InferTaggedBinders]
binds =
  [(Name, TagSig)] -> NameEnv TagSig
forall a. [(Name, a)] -> NameEnv a
mkNameEnv [(Name, TagSig)]
bndr_info
  where
    bndr_info :: [(Name, TagSig)]
bndr_info = (GenStgTopBinding 'InferTaggedBinders -> [(Name, TagSig)])
-> [GenStgTopBinding 'InferTaggedBinders] -> [(Name, TagSig)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap GenStgTopBinding 'InferTaggedBinders -> [(Name, TagSig)]
forall {pass :: StgPass}.
(BinderP pass ~ (Id, TagSig)) =>
GenStgTopBinding pass -> [(Name, TagSig)]
collect [GenStgTopBinding 'InferTaggedBinders]
binds :: [(Name,TagSig)]

    collect :: GenStgTopBinding pass -> [(Name, TagSig)]
collect (StgTopStringLit {}) = []
    collect (StgTopLifted GenStgBinding pass
bnd) =
      case GenStgBinding pass
bnd of
        StgNonRec (Id
id,TagSig
sig) GenStgRhs pass
_rhs
          | TagSig TagInfo
TagDunno <- TagSig
sig -> []
          | Bool
otherwise -> [(Id -> Name
idName Id
id,TagSig
sig)]
        StgRec [(BinderP pass, GenStgRhs pass)]
bnds -> [(BinderP 'InferTaggedBinders, GenStgRhs pass)] -> [(Name, TagSig)]
forall rhs.
[(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
collectRec [(BinderP pass, GenStgRhs pass)]
[(BinderP 'InferTaggedBinders, GenStgRhs pass)]
bnds

    collectRec :: [(BinderP 'InferTaggedBinders, rhs)] -> [(Name,TagSig)]
    collectRec :: forall rhs.
[(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
collectRec [] = []
    collectRec ((BinderP 'InferTaggedBinders, rhs)
bnd:[(BinderP 'InferTaggedBinders, rhs)]
bnds)
      | (BinderP 'InferTaggedBinders
p,rhs
_rhs)  <- (BinderP 'InferTaggedBinders, rhs)
bnd
      , (Id
id,TagSig
sig) <- BinderP 'InferTaggedBinders
p
      , TagSig TagInfo
TagDunno <- TagSig
sig
      = (Id -> Name
idName Id
id,TagSig
sig) (Name, TagSig) -> [(Name, TagSig)] -> [(Name, TagSig)]
forall a. a -> [a] -> [a]
: [(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
forall rhs.
[(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
collectRec [(BinderP 'InferTaggedBinders, rhs)]
bnds
      | Bool
otherwise = [(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
forall rhs.
[(BinderP 'InferTaggedBinders, rhs)] -> [(Name, TagSig)]
collectRec [(BinderP 'InferTaggedBinders, rhs)]
bnds