{-# 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.EnforceEpt ( enforceEpt ) 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.Demand (isDeadEndAppSig)
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.EnforceEpt.Types
import GHC.Stg.EnforceEpt.Rewrite (rewriteTopBinds)
import Data.Maybe
import GHC.Types.Name.Env (mkNameEnv, NameEnv)
import GHC.Driver.DynFlags
import GHC.Utils.Logger
import qualified GHC.Unit.Types

{- Note [Evaluated and Properly Tagged]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
A pointer is Evaluated and Properly Tagged (EPT) when the pointer

  (a) points directly to the value (not to an indirection, and not to a thunk)
  (b) is tagged with the tag corresponding to said value (e.g. constructor tag
      or arity of a function).

A binder is EPT when all the runtime pointers it binds are EPT.

Note that a lifted EPT pointer will never point to a thunk, nor will it be
tagged `000` (meaning "might be a thunk").

See https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/rts/haskell-execution/pointer-tagging
for more information on pointer tagging.

Examples:
* Case binders are always EPT; hence an eval
    case x of x' { __DEFAULT -> ... }
  ensures that x' is EPT even if x was not.
* Data constructor bindings
    let x = Just y in ...
  are EPT: x will point to the heap-allocated constructor closure for (Just y),
  and the tag-bits of the pointer will encode the tag for Just (i.e. `010`).
* In practice, GHC also guarantees that strict fields (and others) are EPT;
  see Note [EPT enforcement].

Caveat:
Currently, the proper tag for builtin *unlifted* data types such as `Array#` is
not `001` but `000`, which is not a proper tag for lifted data.
This means that UnliftedRep is not a proper sub-rep of LiftedRep.
SG thinks it would be good to fix this; see #21792.

Note [EPT enforcement]
~~~~~~~~~~~~~~~~~~~~~~
The goal of EnforceEPT pass is to mark as many binders as possible as EPT
(see Note [Evaluated and Properly Tagged]).
To find more EPT binders, it establishes the following

EPT INVARIANT:
> Any binder of
>   * a strict field (see Note [Strict fields in Core]), or
>   * a CBV argument (see Note [CBV Function Ids])
> is EPT.

(Note that prior to EPT enforcement, this invariant may *not* always be upheld.
An example can be found at the end of this Note.)
This is all to optimise code such as the following:

  data SPair a b = SP !a !b
  case p :: SP Bool Bool of
    SP x y ->
      case x of
        True  -> ...
        False -> ...

We can infer that the strict field x is EPT and hence may safely
omit the code to enter x and the check for the presence of a tag that goes along
with it. However we still branch on the tag as usual to jump to the True or
False case.

Note that for every example involving strict fields we could find a similar
example using CBV functions, e.g.

  $wf x[EPT] y =
    case x of
      True  -> ...
      False -> ...

is the above example translated to use a CBV function $wf.
Note that /any/ strict function can in principle be chosen as a CBV function;
however, we presently only promote worker functions such as $wf to CBV because
we see all its call sites and can use the proper by-value calling convention.
More precisely, with -O0, we guarantee that no CBV functions are visible in
the interface file, so that naïve clients do not need to know how to call CBV
functions. See Note [CBV Function Ids] for more details.

Specification
-------------
EPT enforcement works like implicit type conversions in C, such as from int to
float, only much simpler (no overloaded operations such as +).
For EPT enforcement, the "type system" in question is whether a binder is
statically EPT. We differentiate "EPT binder" from "non-EPT binder", where the
latter means "might be EPT, but we could not prove it so".
In this sense, EPT binders form a subtype of non-EPT binders.
We differentiate two conversion directions:

  * Downcast: EPT binders can be converted into non-EPT binders for free.
  * Upcast: non-EPT binders can be converted into EPT binders by inserting an eval.

The EPT invariant expresses type signatures. In particular, these type
signatures entail two things:

  * A _precondition_: Any binder that is passed as a CBV arg/strict field
    must be EPT (i.e. must have type "EPT binder").
  * A _postcondition_: Any binder of a CBV arg/strict field is EPT.

EPT enforcement is then simply a matter of figuring out where to insert
Upcasts (remember that Downcasts are free).
Since Upcasts (evals!) are not free, it is desirable to insert as few as possible.
To this end, we run a static *EPT analysis*, the purpose of which is to identify
as many EPT binders as possible.
Beyond discovering case binders and value bindings, EPT analysis exploits the
type signatures provided by the EPT invariant, looks inside returned tuples and
does some limited amount of fixpointing.
Afterwards, the *EPT rewriter* inserts the actual evals realising Upcasts.

Implementation
--------------

* EPT analysis is implemented in GHC.Stg.EnforceEpt.inferTags.
  It attaches its result to /binders/, not occurrence sites.
* The EPT rewriter establishes the EPT invariant by inserting evals. That is, if
    (a) a binder x is used to
          * construct a strict field (`SP x y`), or
          * passed as a CBV argument (`$wf x`),
        and
    (b) x was not inferred EPT,
  then the EPT rewriter inserts an eval prior to the call, e.g.
    case x of x' { __ DEFAULT -> SP x' y }.
    case x of x' { __ DEFAULT -> $wf x' }.
  (Recall that the case binder x' is always EPT.)
  This is implemented in GHC.Stg.EnforceEpt.Rewrite.rewriteTopBinds.
  This pass also propagates the EPTness from binders to occurrences.
  It is sound to insert evals on strict fields (Note [Strict fields in Core]),
  and on CBV arguments as well (Note [CBV Function Ids]).
* We also export the EPTness of top level bindings to allow this optimisation
  to work across module boundaries.
  NB: The EPT Invariant *must* be upheld, regardless of the optimisation level;
  hence EPTness is practically part of the internal ABI of a strict data
  constructor or CBV function. Note [CBV Function Ids] contains the details.
* Finally, code generation skips the thunk check when branching on binders that
  are EPT. This is done by `cgExpr`/`cgCase` in the backend.

Evaluation
----------
EPT enforcement can have large impact on spine-strict tree data structure
performance. 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.

Note that EPT enforcement may cause regressions in rare cases.
For example consider this code:

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

When x cannot be inferred EPT, the rewriter transforms to

  foo x = ...
    let c = case x of x' -> StrictJust x'
    in ...

which allocates an additional thunk for `c` that returns the constructor.  Boo!

Note [EPT enforcement lowers strict constructor worker semantics]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In Core, a saturated application of a strict constructor worker evaluates its
strict fields and thus is *not* a value; see Note [Strict fields in Core].
This is also the semantics of strict constructor workers in STG *before* EPT
enforcement (see Note [EPT enforcement])

However, after enforcing the EPT Invariant, all constructor workers can
effectively be lazy. That is, when actually generating code to allocate the
data constructor, the code generator does not need to evaluate the argument;
that has already been done by the EPT pass.

Thus for code-gen reasons (StgToX), all constructor workers are considered lazy
after EPT enforcement.

Note [Why isn't the EPT Invariant enforced during Core passes?]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Recall the definition of the EPT Invariant from Note [EPT enforcement].
Why can't it be established as an invariant right while desugaring to Core?
The reason is that some Core optimisations, such as FloatOut, will drop or delay
evals whenever they think it useful and thus destroy the Invariant.  Example:

  data Set a = Tip | Bin !a (Set a) (Set a)

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 Core optimisation passes
might destroy the EPT Invariant, hence we need to enforce the EPT invariant
*after* passes such as FloatOut.
-}

{- 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 rigorous in dealing with functions.

NB: It turned in #21193 that PAPs get tag zero, so the tag check can't be omitted for functions.
So option two isn't really an option without reworking this anyway.

Note [EPT enforcement debugging]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
There is a flag -dtag-inference-checks which inserts various
compile/runtime checks in order to ensure the EPT 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.

Note [EPT enforcement for interpreted code]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The bytecode interpreter has a different behaviour when it comes
to the tagging of binders in certain situations than the StgToCmm code generator.

a) Tags for let-bindings:

  When compiling a binding for a constructor like `let x = Just True`
  Whether `x` will be properly tagged depends on the backend.
  For the interpreter x points to a BCO which once
  evaluated returns a properly tagged pointer to the heap object.
  In the Cmm backend for the same binding we would allocate the constructor right
  away and x will immediately be represented by a tagged pointer.
  This means for interpreted code we can not assume let bound constructors are
  properly tagged. Hence we distinguish between targeting bytecode and native in
  the analysis.
  We make this differentiation in `mkLetSig` where we simply never assume
  lets are tagged when targeting bytecode.

b) When referencing ids from other modules the Cmm backend will try to put a
   proper tag on these references through various means. When doing analysis we
   usually predict these cases to improve precision of the analysis.
   But to my knowledge the bytecode generator makes no such attempts so we must
   not infer imported bindings as tagged.
   This is handled in GHC.Stg.EnforceEpt.Types.lookupInfo


-}

{- *********************************************************************
*                                                                      *
                         EPT enforcement pass
*                                                                      *
********************************************************************* -}

enforceEpt :: StgPprOpts -> Bool -> Logger -> (GHC.Unit.Types.Module) -> [CgStgTopBinding] -> IO ([TgStgTopBinding], NameEnv TagSig)
enforceEpt :: StgPprOpts
-> Bool
-> Logger
-> Module
-> [CgStgTopBinding]
-> IO ([CgStgTopBinding], NameEnv TagSig)
enforceEpt StgPprOpts
ppr_opts !Bool
for_bytecode Logger
logger Module
this_mod [CgStgTopBinding]
stg_binds = do
    -- pprTraceM "enforceEpt for " (ppr this_mod <> text " bytecode:" <> ppr for_bytecode)
    -- Annotate binders with tag information.
    let (![GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags) = {-# SCC "StgEptInfer" #-}
                                        Bool -> [CgStgTopBinding] -> [GenStgTopBinding 'InferTaggedBinders]
inferTags Bool
for_bytecode [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 StgPprOpts
ppr_opts [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
    us_t <- Char -> IO UniqSupply
mkSplitUniqSupply Char
't'
    let rewritten_binds = {-# SCC "StgEptRewrite" #-} Module
-> UniqSupply
-> [GenStgTopBinding 'InferTaggedBinders]
-> [CgStgTopBinding]
rewriteTopBinds Module
this_mod UniqSupply
us_t [GenStgTopBinding 'InferTaggedBinders]
stg_binds_w_tags :: [TgStgTopBinding]

    return (rewritten_binds,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)

inferTags :: Bool -> [GenStgTopBinding 'CodeGen] -> [GenStgTopBinding 'InferTaggedBinders]
inferTags :: Bool -> [CgStgTopBinding] -> [GenStgTopBinding 'InferTaggedBinders]
inferTags Bool
for_bytecode [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 (Bool -> TagEnv 'CodeGen
initEnv Bool
for_bytecode) [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
         -- It's important that we check for bottoms before all else.
         -- See Note [Bottom functions are TagTagged] and #24806 for why.
         | DmdSig -> Arity -> Bool
isDeadEndAppSig (Id -> DmdSig
idDmdSig Id
fun) ([StgArg] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [StgArg]
args)
         = TagInfo
TagTagged

         | Arity
fun_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
0 -- Unknown arity => Thunk or unknown call
         = TagInfo
TagDunno

         | Just (TagSig TagInfo
res_info) <- IdInfo -> Maybe TagSig
tagSigInfo (HasDebugCallStack => 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

inferTagExpr TagEnv p
env (StgConApp DataCon
con ConstructorNumber
cn [StgArg]
args [[PrimRep]]
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]
-> [[PrimRep]]
-> GenStgExpr 'InferTaggedBinders
forall (pass :: StgPass).
DataCon
-> ConstructorNumber -> [StgArg] -> [[PrimRep]] -> GenStgExpr pass
StgConApp DataCon
con ConstructorNumber
cn [StgArg]
args [[PrimRep]]
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)
  -- Which primops guarantee to return a properly tagged value?
  -- Probably none, and that is the conservative assumption anyway.
  -- (And foreign calls definitely need not make promises.)
  = (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.
HasDebugCallStack =>
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=bndrs', alt_rhs=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 = HasDebugCallStack => 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
      | HasDebugCallStack => 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
out_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
    (TagSig
in_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
    out_sig :: TagSig
out_sig = TagEnv p -> TagSig -> TagSig
forall (p :: StgPass). TagEnv p -> TagSig -> TagSig
mkLetSig TagEnv p
in_env TagSig
in_sig
    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
out_sig)]

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 = 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
         in_bndrs :: [(Id, TagSig)]
in_bndrs = [Id]
in_ids [Id] -> [TagSig] -> [(Id, TagSig)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [TagSig]
in_sigs
         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
         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.
HasDebugCallStack =>
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 =
            let (TagSig
sig_rhs,GenStgRhs 'InferTaggedBinders
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
            in (TagEnv q -> TagSig -> TagSig
forall (p :: StgPass). TagEnv p -> TagSig -> TagSig
mkLetSig TagEnv q
go_env TagSig
sig_rhs, GenStgRhs 'InferTaggedBinders
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
_ Type
_) =
  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.

NB: After the analysis is done we go back to treating bottoming functions as
untagged to ensure they are evaluated as expected in code like:

  case bottom_id of { ...}

-}

-----------------------------
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 Type
typ)
  | 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
-> Type
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> Type
-> GenStgRhs pass
StgRhsClosure XRhsClosure p
XRhsClosure 'InferTaggedBinders
ext CostCentreStack
cc UpdateFlag
upd [(Id, TagSig)]
[BinderP 'InferTaggedBinders]
out_bndrs GenStgExpr 'InferTaggedBinders
body' Type
typ)
  | Bool
otherwise
  = --pprTrace "inferTagRhsClosure" (ppr (_top, _grp_ids, env,info')) $
    (TagInfo -> TagSig
TagSig TagInfo
info', XRhsClosure 'InferTaggedBinders
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'InferTaggedBinders]
-> GenStgExpr 'InferTaggedBinders
-> Type
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> Type
-> GenStgRhs pass
StgRhsClosure XRhsClosure p
XRhsClosure 'InferTaggedBinders
ext CostCentreStack
cc UpdateFlag
upd [(Id, TagSig)]
[BinderP 'InferTaggedBinders]
out_bndrs GenStgExpr 'InferTaggedBinders
body' Type
typ)
  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
_
              | HasDebugCallStack => 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 Type
typ)
-- 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]
-> Type
-> GenStgRhs 'InferTaggedBinders
forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> Type
-> GenStgRhs pass
StgRhsCon CostCentreStack
cc DataCon
con ConstructorNumber
cn [StgTickish]
ticks [StgArg]
args Type
typ)

-- Adjust let semantics to the targeted backend.
-- See Note [EPT enforcement for interpreted code]
mkLetSig :: TagEnv p -> TagSig -> TagSig
mkLetSig :: forall (p :: StgPass). TagEnv p -> TagSig -> TagSig
mkLetSig TagEnv p
env TagSig
in_sig
  | Bool
for_bytecode = TagInfo -> TagSig
TagSig TagInfo
TagDunno
  | Bool
otherwise = TagSig
in_sig
  where
    for_bytecode :: Bool
for_bytecode = TagEnv p -> Bool
forall (p :: StgPass). TagEnv p -> Bool
te_bytecode TagEnv p
env

{- 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 scrutinize 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. HasDebugCallStack => String -> [a] -> [b] -> [(a, b)]
zipEqual String
"inferTagRhs" [StgArg]
args (HasDebugCallStack => 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