{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}

module RepType
  (
    -- * Code generator views onto Types
    UnaryType, NvUnaryType, isNvUnaryType,
    unwrapType,

    -- * Predicates on types
    isVoidTy,

    -- * Type representation for the code generator
    typePrimRep, typePrimRep1,
    runtimeRepPrimRep, typePrimRepArgs,
    PrimRep(..), primRepToType,
    countFunRepArgs, countConRepArgs, tyConPrimRep, tyConPrimRep1,

    -- * Unboxed sum representation type
    ubxSumRepType, layoutUbxSum, typeSlotTy, SlotTy (..),
    slotPrimRep, primRepSlot
  ) where

#include "HsVersions.h"

import GhcPrelude

import BasicTypes (Arity, RepArity)
import DataCon
import Outputable
import PrelNames
import Coercion
import TyCon
import TyCoRep
import Type
import Util
import TysPrim
import {-# SOURCE #-} TysWiredIn ( anyTypeOfKind )

import Data.List (sort)
import qualified Data.IntSet as IS

{- **********************************************************************
*                                                                       *
                Representation types
*                                                                       *
********************************************************************** -}

type NvUnaryType = Type
type UnaryType   = Type
     -- Both are always a value type; i.e. its kind is TYPE rr
     -- for some rr; moreover the rr is never a variable.
     --
     --   NvUnaryType : never an unboxed tuple or sum, or void
     --
     --   UnaryType   : never an unboxed tuple or sum;
     --                 can be Void# or (# #)

isNvUnaryType :: Type -> Bool
isNvUnaryType ty
  | [_] <- typePrimRep ty
  = True
  | otherwise
  = False

-- INVARIANT: the result list is never empty.
typePrimRepArgs :: Type -> [PrimRep]
typePrimRepArgs ty
  | [] <- reps
  = [VoidRep]
  | otherwise
  = reps
  where
    reps = typePrimRep ty

-- | Gets rid of the stuff that prevents us from understanding the
-- runtime representation of a type. Including:
--   1. Casts
--   2. Newtypes
--   3. Foralls
--   4. Synonyms
-- But not type/data families, because we don't have the envs to hand.
unwrapType :: Type -> Type
unwrapType ty
  | Just (_, unwrapped)
      <- topNormaliseTypeX stepper mappend inner_ty
  = unwrapped
  | otherwise
  = inner_ty
  where
    inner_ty = go ty

    go t | Just t' <- coreView t = go t'
    go (ForAllTy _ t)            = go t
    go (CastTy t _)              = go t
    go t                         = t

     -- cf. Coercion.unwrapNewTypeStepper
    stepper rec_nts tc tys
      | Just (ty', _) <- instNewTyCon_maybe tc tys
      = case checkRecTc rec_nts tc of
          Just rec_nts' -> NS_Step rec_nts' (go ty') ()
          Nothing       -> NS_Abort   -- infinite newtypes
      | otherwise
      = NS_Done

countFunRepArgs :: Arity -> Type -> RepArity
countFunRepArgs 0 _
  = 0
countFunRepArgs n ty
  | FunTy _ arg res <- unwrapType ty
  = length (typePrimRepArgs arg) + countFunRepArgs (n - 1) res
  | otherwise
  = pprPanic "countFunRepArgs: arity greater than type can handle" (ppr (n, ty, typePrimRep ty))

countConRepArgs :: DataCon -> RepArity
countConRepArgs dc = go (dataConRepArity dc) (dataConRepType dc)
  where
    go :: Arity -> Type -> RepArity
    go 0 _
      = 0
    go n ty
      | FunTy _ arg res <- unwrapType ty
      = length (typePrimRep arg) + go (n - 1) res
      | otherwise
      = pprPanic "countConRepArgs: arity greater than type can handle" (ppr (n, ty, typePrimRep ty))

-- | True if the type has zero width.
isVoidTy :: Type -> Bool
isVoidTy = null . typePrimRep

{- **********************************************************************
*                                                                       *
                Unboxed sums
 See Note [Translating unboxed sums to unboxed tuples] in UnariseStg.hs
*                                                                       *
********************************************************************** -}

type SortedSlotTys = [SlotTy]

-- | Given the arguments of a sum type constructor application,
--   return the unboxed sum rep type.
--
-- E.g.
--
--   (# Int# | Maybe Int | (# Int#, Float# #) #)
--
-- We call `ubxSumRepType [ [IntRep], [LiftedRep], [IntRep, FloatRep] ]`,
-- which returns [WordSlot, PtrSlot, WordSlot, FloatSlot]
--
-- INVARIANT: Result slots are sorted (via Ord SlotTy), except that at the head
-- of the list we have the slot for the tag.
ubxSumRepType :: [[PrimRep]] -> [SlotTy]
ubxSumRepType constrs0
  -- These first two cases never classify an actual unboxed sum, which always
  -- has at least two disjuncts. But it could happen if a user writes, e.g.,
  -- forall (a :: TYPE (SumRep [IntRep])). ...
  -- which could never be instantiated. We still don't want to panic.
  | constrs0 `lengthLessThan` 2
  = [WordSlot]

  | otherwise
  = let
      combine_alts :: [SortedSlotTys]  -- slots of constructors
                   -> SortedSlotTys    -- final slots
      combine_alts constrs = foldl' merge [] constrs

      merge :: SortedSlotTys -> SortedSlotTys -> SortedSlotTys
      merge existing_slots []
        = existing_slots
      merge [] needed_slots
        = needed_slots
      merge (es : ess) (s : ss)
        | Just s' <- s `fitsIn` es
        = -- found a slot, use it
          s' : merge ess ss
        | s < es
        = -- we need a new slot and this is the right place for it
          s : merge (es : ess) ss
        | otherwise
        = -- keep searching for a slot
          es : merge ess (s : ss)

      -- Nesting unboxed tuples and sums is OK, so we need to flatten first.
      rep :: [PrimRep] -> SortedSlotTys
      rep ty = sort (map primRepSlot ty)

      sumRep = WordSlot : combine_alts (map rep constrs0)
               -- WordSlot: for the tag of the sum
    in
      sumRep

layoutUbxSum :: SortedSlotTys -- Layout of sum. Does not include tag.
                              -- We assume that they are in increasing order
             -> [SlotTy]      -- Slot types of things we want to map to locations in the
                              -- sum layout
             -> [Int]         -- Where to map 'things' in the sum layout
layoutUbxSum sum_slots0 arg_slots0 =
    go arg_slots0 IS.empty
  where
    go :: [SlotTy] -> IS.IntSet -> [Int]
    go [] _
      = []
    go (arg : args) used
      = let slot_idx = findSlot arg 0 sum_slots0 used
         in slot_idx : go args (IS.insert slot_idx used)

    findSlot :: SlotTy -> Int -> SortedSlotTys -> IS.IntSet -> Int
    findSlot arg slot_idx (slot : slots) useds
      | not (IS.member slot_idx useds)
      , Just slot == arg `fitsIn` slot
      = slot_idx
      | otherwise
      = findSlot arg (slot_idx + 1) slots useds
    findSlot _ _ [] _
      = pprPanic "findSlot" (text "Can't find slot" $$ ppr sum_slots0 $$ ppr arg_slots0)

--------------------------------------------------------------------------------

-- We have 3 kinds of slots:
--
--   - Pointer slot: Only shared between actual pointers to Haskell heap (i.e.
--     boxed objects)
--
--   - Word slots: Shared between IntRep, WordRep, Int64Rep, Word64Rep, AddrRep.
--
--   - Float slots: Shared between floating point types.
--
--   - Void slots: Shared between void types. Not used in sums.
--
-- TODO(michalt): We should probably introduce `SlotTy`s for 8-/16-/32-bit
-- values, so that we can pack things more tightly.
data SlotTy = PtrSlot | WordSlot | Word64Slot | FloatSlot | DoubleSlot
  deriving (Eq, Ord)
    -- Constructor order is important! If slot A could fit into slot B
    -- then slot A must occur first.  E.g.  FloatSlot before DoubleSlot
    --
    -- We are assuming that WordSlot is smaller than or equal to Word64Slot
    -- (would not be true on a 128-bit machine)

instance Outputable SlotTy where
  ppr PtrSlot    = text "PtrSlot"
  ppr Word64Slot = text "Word64Slot"
  ppr WordSlot   = text "WordSlot"
  ppr DoubleSlot = text "DoubleSlot"
  ppr FloatSlot  = text "FloatSlot"

typeSlotTy :: UnaryType -> Maybe SlotTy
typeSlotTy ty
  | isVoidTy ty
  = Nothing
  | otherwise
  = Just (primRepSlot (typePrimRep1 ty))

primRepSlot :: PrimRep -> SlotTy
primRepSlot VoidRep     = pprPanic "primRepSlot" (text "No slot for VoidRep")
primRepSlot LiftedRep   = PtrSlot
primRepSlot UnliftedRep = PtrSlot
primRepSlot IntRep      = WordSlot
primRepSlot Int8Rep     = WordSlot
primRepSlot Int16Rep    = WordSlot
primRepSlot Int64Rep    = Word64Slot
primRepSlot WordRep     = WordSlot
primRepSlot Word8Rep    = WordSlot
primRepSlot Word16Rep   = WordSlot
primRepSlot Word64Rep   = Word64Slot
primRepSlot AddrRep     = WordSlot
primRepSlot FloatRep    = FloatSlot
primRepSlot DoubleRep   = DoubleSlot
primRepSlot VecRep{}    = pprPanic "primRepSlot" (text "No slot for VecRep")

slotPrimRep :: SlotTy -> PrimRep
slotPrimRep PtrSlot     = LiftedRep   -- choice between lifted & unlifted seems arbitrary
slotPrimRep Word64Slot  = Word64Rep
slotPrimRep WordSlot    = WordRep
slotPrimRep DoubleSlot  = DoubleRep
slotPrimRep FloatSlot   = FloatRep

-- | Returns the bigger type if one fits into the other. (commutative)
fitsIn :: SlotTy -> SlotTy -> Maybe SlotTy
fitsIn ty1 ty2
  | isWordSlot ty1 && isWordSlot ty2
  = Just (max ty1 ty2)
  | isFloatSlot ty1 && isFloatSlot ty2
  = Just (max ty1 ty2)
  | isPtrSlot ty1 && isPtrSlot ty2
  = Just PtrSlot
  | otherwise
  = Nothing
  where
    isPtrSlot PtrSlot = True
    isPtrSlot _       = False

    isWordSlot Word64Slot = True
    isWordSlot WordSlot   = True
    isWordSlot _          = False

    isFloatSlot DoubleSlot = True
    isFloatSlot FloatSlot  = True
    isFloatSlot _          = False

{- **********************************************************************
*                                                                       *
                   PrimRep
*                                                                       *
********************************************************************** -}

-- | Discovers the primitive representation of a 'Type'. Returns
-- a list of 'PrimRep': it's a list because of the possibility of
-- no runtime representation (void) or multiple (unboxed tuple/sum)
typePrimRep :: HasDebugCallStack => Type -> [PrimRep]
typePrimRep ty = kindPrimRep (text "typePrimRep" <+>
                              parens (ppr ty <+> dcolon <+> ppr (typeKind ty)))
                             (typeKind ty)

-- | Like 'typePrimRep', but assumes that there is precisely one 'PrimRep' output;
-- an empty list of PrimReps becomes a VoidRep
typePrimRep1 :: HasDebugCallStack => UnaryType -> PrimRep
typePrimRep1 ty = case typePrimRep ty of
  []    -> VoidRep
  [rep] -> rep
  _     -> pprPanic "typePrimRep1" (ppr ty $$ ppr (typePrimRep ty))

-- | Find the runtime representation of a 'TyCon'. Defined here to
-- avoid module loops. Returns a list of the register shapes necessary.
tyConPrimRep :: HasDebugCallStack => TyCon -> [PrimRep]
tyConPrimRep tc
  = kindPrimRep (text "kindRep tc" <+> ppr tc $$ ppr res_kind)
                res_kind
  where
    res_kind = tyConResKind tc

-- | Like 'tyConPrimRep', but assumed that there is precisely zero or
-- one 'PrimRep' output
tyConPrimRep1 :: HasDebugCallStack => TyCon -> PrimRep
tyConPrimRep1 tc = case tyConPrimRep tc of
  []    -> VoidRep
  [rep] -> rep
  _     -> pprPanic "tyConPrimRep1" (ppr tc $$ ppr (tyConPrimRep tc))

-- | Take a kind (of shape @TYPE rr@) and produce the 'PrimRep's
-- of values of types of this kind.
kindPrimRep :: HasDebugCallStack => SDoc -> Kind -> [PrimRep]
kindPrimRep doc ki
  | Just ki' <- coreView ki
  = kindPrimRep doc ki'
kindPrimRep doc (TyConApp typ [runtime_rep])
  = ASSERT( typ `hasKey` tYPETyConKey )
    runtimeRepPrimRep doc runtime_rep
kindPrimRep doc ki
  = pprPanic "kindPrimRep" (ppr ki $$ doc)

-- | Take a type of kind RuntimeRep and extract the list of 'PrimRep' that
-- it encodes.
runtimeRepPrimRep :: HasDebugCallStack => SDoc -> Type -> [PrimRep]
runtimeRepPrimRep doc rr_ty
  | Just rr_ty' <- coreView rr_ty
  = runtimeRepPrimRep doc rr_ty'
  | TyConApp rr_dc args <- rr_ty
  , RuntimeRep fun <- tyConRuntimeRepInfo rr_dc
  = fun args
  | otherwise
  = pprPanic "runtimeRepPrimRep" (doc $$ ppr rr_ty)

-- | Convert a PrimRep back to a Type. Used only in the unariser to give types
-- to fresh Ids. Really, only the type's representation matters.
primRepToType :: PrimRep -> Type
primRepToType = anyTypeOfKind . tYPE . primRepToRuntimeRep


{-
(c) The GRASP/AQUA Project, Glasgow University, 1993-1998

\section[SimplStg]{Driver for simplifying @STG@ programs}
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

module SimplStg ( stg2stg ) where

#include "HsVersions.h"

import GhcPrelude

import StgSyn

import StgLint          ( lintStgTopBindings )
import StgStats         ( showStgStats )
import UnariseStg       ( unarise )
import StgCse           ( stgCse )
import StgLiftLams      ( stgLiftLams )
import Module           ( Module )

import DynFlags
import ErrUtils
import UniqSupply
import Outputable
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.State.Strict

newtype StgM a = StgM { _unStgM :: StateT UniqSupply IO a }
  deriving (Functor, Applicative, Monad, MonadIO)

instance MonadUnique StgM where
  getUniqueSupplyM = StgM (state splitUniqSupply)
  getUniqueM = StgM (state takeUniqFromSupply)

runStgM :: UniqSupply -> StgM a -> IO a
runStgM us (StgM m) = evalStateT m us

stg2stg :: DynFlags                  -- includes spec of what stg-to-stg passes to do
        -> Module                    -- module being compiled
        -> [StgTopBinding]           -- input program
        -> IO [StgTopBinding]        -- output program

stg2stg dflags this_mod binds
  = do  { showPass dflags "Stg2Stg"
        ; us <- mkSplitUniqSupply 'g'

        -- Do the main business!
        ; binds' <- runStgM us $
            foldM do_stg_pass binds (getStgToDo dflags)

        ; dump_when Opt_D_dump_stg "STG syntax:" binds'

        ; return binds'
   }

  where
    stg_linter what
      | gopt Opt_DoStgLinting dflags
      = lintStgTopBindings dflags this_mod what
      | otherwise
      = \ _whodunnit _binds -> return ()

    -------------------------------------------
    do_stg_pass :: [StgTopBinding] -> StgToDo -> StgM [StgTopBinding]
    do_stg_pass binds to_do
      = case to_do of
          StgDoNothing ->
            return binds

          StgStats ->
            trace (showStgStats binds) (return binds)

          StgCSE -> do
            let binds' = {-# SCC "StgCse" #-} stgCse binds
            end_pass "StgCse" binds'

          StgLiftLams -> do
            us <- getUniqueSupplyM
            let binds' = {-# SCC "StgLiftLams" #-} stgLiftLams dflags us binds
            end_pass "StgLiftLams" binds'

          StgUnarise -> do
            liftIO (dump_when Opt_D_dump_stg "Pre unarise:" binds)
            us <- getUniqueSupplyM
            liftIO (stg_linter False "Pre-unarise" binds)
            let binds' = unarise us binds
            liftIO (stg_linter True "Unarise" binds')
            return binds'

    dump_when flag header binds
      = dumpIfSet_dyn dflags flag header (pprStgTopBindings binds)

    end_pass what binds2
      = liftIO $ do -- report verbosely, if required
          dumpIfSet_dyn dflags Opt_D_verbose_stg2stg what
            (vcat (map ppr binds2))
          stg_linter False what binds2
          return binds2

-- -----------------------------------------------------------------------------
-- StgToDo:  abstraction of stg-to-stg passes to run.

-- | Optional Stg-to-Stg passes.
data StgToDo
  = StgCSE
  -- ^ Common subexpression elimination
  | StgLiftLams
  -- ^ Lambda lifting closure variables, trading stack/register allocation for
  -- heap allocation
  | StgStats
  | StgUnarise
  -- ^ Mandatory unarise pass, desugaring unboxed tuple and sum binders
  | StgDoNothing
  -- ^ Useful for building up 'getStgToDo'
  deriving Eq

-- | Which Stg-to-Stg passes to run. Depends on flags, ways etc.
getStgToDo :: DynFlags -> [StgToDo]
getStgToDo dflags =
  filter (/= StgDoNothing)
    [ mandatory StgUnarise
    -- Important that unarisation comes first
    -- See Note [StgCse after unarisation] in StgCse
    , optional Opt_StgCSE StgCSE
    , optional Opt_StgLiftLams StgLiftLams
    , optional Opt_StgStats StgStats
    ] where
      optional opt = runWhen (gopt opt dflags)
      mandatory = id

runWhen :: Bool -> StgToDo -> StgToDo
runWhen True todo = todo
runWhen _    _    = StgDoNothing


{-# LANGUAGE TypeFamilies #-}

{-|
Note [CSE for Stg]
~~~~~~~~~~~~~~~~~~
This module implements a simple common subexpression elimination pass for STG.
This is useful because there are expressions that we want to common up (because
they are operationally equivalent), but that we cannot common up in Core, because
their types differ.
This was originally reported as #9291.

There are two types of common code occurrences that we aim for, see
note [Case 1: CSEing allocated closures] and
note [Case 2: CSEing case binders] below.

Note [Case 1: CSEing allocated closures]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The first kind of CSE opportunity we aim for is generated by this Haskell code:

    bar :: a -> (Either Int a, Either Bool a)
    bar x = (Right x, Right x)

which produces this Core:

    bar :: forall a. a -> (Either Int a, Either Bool a)
    bar @a x = (Right @Int @a x, Right @Bool @a x)

where the two components of the tuple are different terms, and cannot be
commoned up (easily). On the STG level we have

    bar [x] = let c1 = Right [x]
                  c2 = Right [x]
              in (c1,c2)

and now it is obvious that we can write

    bar [x] = let c1 = Right [x]
              in (c1,c1)

instead.

Note [Case 2: CSEing case binders]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The second kind of CSE opportunity we aim for is more interesting, and
came up in #9291 and #5344: The Haskell code

    foo :: Either Int a -> Either Bool a
    foo (Right x) = Right x
    foo _         = Left False

produces this Core

    foo :: forall a. Either Int a -> Either Bool a
    foo @a e = case e of b { Left n -> â€¦
                           , Right x -> Right @Bool @a x }

where we cannot CSE `Right @Bool @a x` with the case binder `b` as they have
different types. But in STG we have

    foo [e] = case e of b { Left [n] -> â€¦
                          , Right [x] -> Right [x] }

and nothing stops us from transforming that to

    foo [e] = case e of b { Left [n] -> â€¦
                          , Right [x] -> b}

Note [StgCse after unarisation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider two unboxed sum terms:

    (# 1 | #) :: (# Int | Int# #)
    (# 1 | #) :: (# Int | Int  #)

These two terms are not equal as they unarise to different unboxed
tuples. However if we run StgCse before Unarise, it'll think the two
terms (# 1 | #) are equal, and replace one of these with a binder to
the other. That's bad -- Trac #15300.

Solution: do unarise first.

-}

module StgCse (stgCse) where

import GhcPrelude

import DataCon
import Id
import StgSyn
import Outputable
import VarEnv
import CoreSyn (AltCon(..))
import Data.List (mapAccumL)
import Data.Maybe (fromMaybe)
import CoreMap
import NameEnv
import Control.Monad( (>=>) )

--------------
-- The Trie --
--------------

-- A lookup trie for data constructor applications, i.e.
-- keys of type `(DataCon, [StgArg])`, following the patterns in TrieMap.

data StgArgMap a = SAM
    { sam_var :: DVarEnv a
    , sam_lit :: LiteralMap a
    }

instance TrieMap StgArgMap where
    type Key StgArgMap = StgArg
    emptyTM  = SAM { sam_var = emptyTM
                   , sam_lit = emptyTM }
    lookupTM (StgVarArg var) = sam_var >.> lkDFreeVar var
    lookupTM (StgLitArg lit) = sam_lit >.> lookupTM lit
    alterTM  (StgVarArg var) f m = m { sam_var = sam_var m |> xtDFreeVar var f }
    alterTM  (StgLitArg lit) f m = m { sam_lit = sam_lit m |> alterTM lit f }
    foldTM k m = foldTM k (sam_var m) . foldTM k (sam_lit m)
    mapTM f (SAM {sam_var = varm, sam_lit = litm}) =
        SAM { sam_var = mapTM f varm, sam_lit = mapTM f litm }

newtype ConAppMap a = CAM { un_cam :: DNameEnv (ListMap StgArgMap a) }

instance TrieMap ConAppMap where
    type Key ConAppMap = (DataCon, [StgArg])
    emptyTM  = CAM emptyTM
    lookupTM (dataCon, args) = un_cam >.> lkDNamed dataCon >=> lookupTM args
    alterTM  (dataCon, args) f m =
        m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
    foldTM k = un_cam >.> foldTM (foldTM k)
    mapTM f  = un_cam >.> mapTM (mapTM f) >.> CAM

-----------------
-- The CSE Env --
-----------------

-- | The CSE environment. See note [CseEnv Example]
data CseEnv = CseEnv
    { ce_conAppMap :: ConAppMap OutId
        -- ^ The main component of the environment is the trie that maps
        --   data constructor applications (with their `OutId` arguments)
        --   to an in-scope name that can be used instead.
        --   This name is always either a let-bound variable or a case binder.
    , ce_subst     :: IdEnv OutId
        -- ^ This substitution is applied to the code as we traverse it.
        --   Entries have one of two reasons:
        --
        --   * The input might have shadowing (see Note [Shadowing]), so we have
        --     to rename some binders as we traverse the tree.
        --   * If we remove `let x = Con z` because  `let y = Con z` is in scope,
        --     we note this here as x â†¦ y.
    , ce_bndrMap     :: IdEnv OutId
        -- ^ If we come across a case expression case x as b of â€¦ with a trivial
        --   binder, we add b â†¦ x to this.
        --   This map is *only* used when looking something up in the ce_conAppMap.
        --   See Note [Trivial case scrutinee]
    , ce_in_scope  :: InScopeSet
        -- ^ The third component is an in-scope set, to rename away any
        --   shadowing binders
    }

{-|
Note [CseEnv Example]
~~~~~~~~~~~~~~~~~~~~~
The following tables shows how the CseEnvironment changes as code is traversed,
as well as the changes to that code.

  InExpr                         OutExpr
     conAppMap                   subst          in_scope
  â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€
  -- empty                       {}             {}
  case â€¦ as a of {Con x y ->     case â€¦ as a of {Con x y ->
  -- Con x y â†¦ a                 {}             {a,x,y}
  let b = Con x y                (removed)
  -- Con x y â†¦ a                 bâ†¦a            {a,x,y,b}
  let c = Bar a                  let c = Bar a
  -- Con x y â†¦ a, Bar a â†¦ c      bâ†¦a            {a,x,y,b,c}
  let c = some expression        let c' = some expression
  -- Con x y â†¦ a, Bar a â†¦ c      bâ†¦a, câ†¦c',     {a,x,y,b,c,c'}
  let d = Bar b                  (removed)
  -- Con x y â†¦ a, Bar a â†¦ c      bâ†¦a, câ†¦c', dâ†¦c {a,x,y,b,c,c',d}
  (a, b, c d)                    (a, a, c' c)
-}

initEnv :: InScopeSet -> CseEnv
initEnv in_scope = CseEnv
    { ce_conAppMap = emptyTM
    , ce_subst     = emptyVarEnv
    , ce_bndrMap   = emptyVarEnv
    , ce_in_scope  = in_scope
    }

envLookup :: DataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
envLookup dataCon args env = lookupTM (dataCon, args') (ce_conAppMap env)
  where args' = map go args -- See Note [Trivial case scrutinee]
        go (StgVarArg v  ) = StgVarArg (fromMaybe v $ lookupVarEnv (ce_bndrMap env) v)
        go (StgLitArg lit) = StgLitArg lit

addDataCon :: OutId -> DataCon -> [OutStgArg] -> CseEnv -> CseEnv
-- do not bother with nullary data constructors, they are static anyways
addDataCon _ _ [] env = env
addDataCon bndr dataCon args env = env { ce_conAppMap = new_env }
  where
    new_env = insertTM (dataCon, args) bndr (ce_conAppMap env)

forgetCse :: CseEnv -> CseEnv
forgetCse env = env { ce_conAppMap = emptyTM }
    -- See note [Free variables of an StgClosure]

addSubst :: OutId -> OutId -> CseEnv -> CseEnv
addSubst from to env
    = env { ce_subst = extendVarEnv (ce_subst env) from to }

addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
addTrivCaseBndr from to env
    = env { ce_bndrMap = extendVarEnv (ce_bndrMap env) from to }

substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
substArgs env = map (substArg env)

substArg :: CseEnv -> InStgArg -> OutStgArg
substArg env (StgVarArg from) = StgVarArg (substVar env from)
substArg _   (StgLitArg lit)  = StgLitArg lit

substVar :: CseEnv -> InId -> OutId
substVar env id = fromMaybe id $ lookupVarEnv (ce_subst env) id

-- Functions to enter binders

-- This is much simpler than the equivalent code in CoreSubst:
--  * We do not substitute type variables, and
--  * There is nothing relevant in IdInfo at this stage
--    that needs substitutions.
-- Therefore, no special treatment for a recursive group is required.

substBndr :: CseEnv -> InId -> (CseEnv, OutId)
substBndr env old_id
  = (new_env, new_id)
  where
    new_id = uniqAway (ce_in_scope env) old_id
    no_change = new_id == old_id
    env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
    new_env | no_change = env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }
            | otherwise = env'

substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
substBndrs env bndrs = mapAccumL substBndr env bndrs

substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
substPairs env bndrs = mapAccumL go env bndrs
  where go env (id, x) = let (env', id') = substBndr env id
                         in (env', (id', x))

-- Main entry point

stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
stgCse binds = snd $ mapAccumL stgCseTopLvl emptyInScopeSet binds

-- Top level bindings.
--
-- We do not CSE these, as top-level closures are allocated statically anyways.
-- Also, they might be exported.
-- But we still have to collect the set of in-scope variables, otherwise
-- uniqAway might shadow a top-level closure.

stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
stgCseTopLvl in_scope t@(StgTopStringLit _ _) = (in_scope, t)
stgCseTopLvl in_scope (StgTopLifted (StgNonRec bndr rhs))
    = (in_scope'
      , StgTopLifted (StgNonRec bndr (stgCseTopLvlRhs in_scope rhs)))
  where in_scope' = in_scope `extendInScopeSet` bndr

stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
    = ( in_scope'
      , StgTopLifted (StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ]))
  where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]

stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
stgCseTopLvlRhs in_scope (StgRhsClosure ext ccs upd args body)
    = let body' = stgCseExpr (initEnv in_scope) body
      in  StgRhsClosure ext ccs upd args body'
stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
    = StgRhsCon ccs dataCon args

------------------------------
-- The actual AST traversal --
------------------------------

-- Trivial cases
stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
stgCseExpr env (StgApp fun args)
    = StgApp fun' args'
  where fun' = substVar env fun
        args' = substArgs env args
stgCseExpr _ (StgLit lit)
    = StgLit lit
stgCseExpr env (StgOpApp op args tys)
    = StgOpApp op args' tys
  where args' = substArgs env args
stgCseExpr _ (StgLam _ _)
    = pprPanic "stgCseExp" (text "StgLam")
stgCseExpr env (StgTick tick body)
    = let body' = stgCseExpr env body
      in StgTick tick body'
stgCseExpr env (StgCase scrut bndr ty alts)
    = mkStgCase scrut' bndr' ty alts'
  where
    scrut' = stgCseExpr env scrut
    (env1, bndr') = substBndr env bndr
    env2 | StgApp trivial_scrut [] <- scrut' = addTrivCaseBndr bndr trivial_scrut env1
                 -- See Note [Trivial case scrutinee]
         | otherwise                         = env1
    alts' = map (stgCseAlt env2 ty bndr') alts

-- A constructor application.
-- To be removed by a variable use when found in the CSE environment
stgCseExpr env (StgConApp dataCon args tys)
    | Just bndr' <- envLookup dataCon args' env
    = StgApp bndr' []
    | otherwise
    = StgConApp dataCon args' tys
  where args' = substArgs env args

-- Let bindings
-- The binding might be removed due to CSE (we do not want trivial bindings on
-- the STG level), so use the smart constructor `mkStgLet` to remove the binding
-- if empty.
stgCseExpr env (StgLet ext binds body)
    = let (binds', env') = stgCseBind env binds
          body' = stgCseExpr env' body
      in mkStgLet (StgLet ext) binds' body'
stgCseExpr env (StgLetNoEscape ext binds body)
    = let (binds', env') = stgCseBind env binds
          body' = stgCseExpr env' body
      in mkStgLet (StgLetNoEscape ext) binds' body'

-- Case alternatives
-- Extend the CSE environment
stgCseAlt :: CseEnv -> AltType -> OutId -> InStgAlt -> OutStgAlt
stgCseAlt env ty case_bndr (DataAlt dataCon, args, rhs)
    = let (env1, args') = substBndrs env args
          env2
            -- To avoid dealing with unboxed sums StgCse runs after unarise and
            -- should maintain invariants listed in Note [Post-unarisation
            -- invariants]. One of the invariants is that some binders are not
            -- used (unboxed tuple case binders) which is what we check with
            -- `stgCaseBndrInScope` here. If the case binder is not in scope we
            -- don't add it to the CSE env. See also #15300.
            | stgCaseBndrInScope ty True -- CSE runs after unarise
            = addDataCon case_bndr dataCon (map StgVarArg args') env1
            | otherwise
            = env1
            -- see note [Case 2: CSEing case binders]
          rhs' = stgCseExpr env2 rhs
      in (DataAlt dataCon, args', rhs')
stgCseAlt env _ _ (altCon, args, rhs)
    = let (env1, args') = substBndrs env args
          rhs' = stgCseExpr env1 rhs
      in (altCon, args', rhs')

-- Bindings
stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
stgCseBind env (StgNonRec b e)
    = let (env1, b') = substBndr env b
      in case stgCseRhs env1 b' e of
        (Nothing,      env2) -> (Nothing,                env2)
        (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
stgCseBind env (StgRec pairs)
    = let (env1, pairs1) = substPairs env pairs
      in case stgCsePairs env1 pairs1 of
        ([],     env2) -> (Nothing, env2)
        (pairs2, env2) -> (Just (StgRec pairs2), env2)

stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
stgCsePairs env [] = ([], env)
stgCsePairs env0 ((b,e):pairs)
  = let (pairMB, env1) = stgCseRhs env0 b e
        (pairs', env2) = stgCsePairs env1 pairs
    in (pairMB `mbCons` pairs', env2)
  where
    mbCons = maybe id (:)

-- The RHS of a binding.
-- If it is a constructor application, either short-cut it or extend the environment
stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
stgCseRhs env bndr (StgRhsCon ccs dataCon args)
    | Just other_bndr <- envLookup dataCon args' env
    = let env' = addSubst bndr other_bndr env
      in (Nothing, env')
    | otherwise
    = let env' = addDataCon bndr dataCon args' env
            -- see note [Case 1: CSEing allocated closures]
          pair = (bndr, StgRhsCon ccs dataCon args')
      in (Just pair, env')
  where args' = substArgs env args
stgCseRhs env bndr (StgRhsClosure ext ccs upd args body)
    = let (env1, args') = substBndrs env args
          env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
          body' = stgCseExpr env2 body
      in (Just (substVar env bndr, StgRhsClosure ext ccs upd args' body'), env)

mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
mkStgCase scrut bndr ty alts | all isBndr alts = scrut
                             | otherwise       = StgCase scrut bndr ty alts

  where
    -- see Note [All alternatives are the binder]
    isBndr (_, _, StgApp f []) = f == bndr
    isBndr _                   = False

-- Utilities

-- | This function short-cuts let-bindings that are now obsolete
mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
mkStgLet _      Nothing      body = body
mkStgLet stgLet (Just binds) body = stgLet binds body

{-
Note [All alternatives are the binder]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

When all alternatives simply refer to the case binder, then we do not have
to bother with the case expression at all (#13588). CoreSTG does this as well,
but sometimes, types get into the way:

    newtype T = MkT Int
    f :: (Int, Int) -> (T, Int)
    f (x, y) = (MkT x, y)

Core cannot just turn this into

    f p = p

as this would not be well-typed. But to STG, where MkT is no longer in the way,
we can.

Note [Trivial case scrutinee]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to be able to handle nested reconstruction of constructors as in

    nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
    nested (Right (Right v)) = Right (Right v)
    nested _ = Left True

So if we come across

    case x of r1
      Right a -> case a of r2
              Right b -> let v = Right b
                         in Right v

we first replace v with r2. Next we want to replace Right r2 with r1. But the
ce_conAppMap contains Right a!

Therefore, we add r1 â†¦ x to ce_bndrMap when analysing the outer case, and use
this substitution before looking Right r2 up in ce_conAppMap, and everything
works out.

Note [Free variables of an StgClosure]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
StgClosures (function and thunks) have an explicit list of free variables:

foo [x] =
    let not_a_free_var = Left [x]
    let a_free_var = Right [x]
    let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
    in closure

If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
then the list of free variables would be wrong, so for now, we do not CSE
across such a closure, simply because I (Joachim) was not sure about possible
knock-on effects. If deemed safe and worth the slight code complication of
re-calculating this list during or after this pass, this can surely be done.
-}


-- | Implements a selective lambda lifter, running late in the optimisation
-- pipeline.
--
-- The transformation itself is implemented in "StgLiftLams.Transformation".
-- If you are interested in the cost model that is employed to decide whether
-- to lift a binding or not, look at "StgLiftLams.Analysis".
-- "StgLiftLams.LiftM" contains the transformation monad that hides away some
-- plumbing of the transformation.
module StgLiftLams (
    -- * Late lambda lifting in STG
    -- $note
    Transformation.stgLiftLams
  ) where

import qualified StgLiftLams.Transformation as Transformation

-- Note [Late lambda lifting in STG]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- $note
-- See also the <https://ghc.haskell.org/trac/ghc/wiki/LateLamLift wiki page>
-- and Trac #9476.
--
-- The basic idea behind lambda lifting is to turn locally defined functions
-- into top-level functions. Free variables are then passed as additional
-- arguments at *call sites* instead of having a closure allocated for them at
-- *definition site*. Example:
--
-- @
--    let x = ...; y = ... in
--    let f = {x y} \a -> a + x + y in
--    let g = {f x} \b -> f b + x in
--    g 5
-- @
--
-- Lambda lifting @f@ would
--
--   1. Turn @f@'s free variables into formal parameters
--   2. Update @f@'s call site within @g@ to @f x y b@
--   3. Update @g@'s closure: Add @y@ as an additional free variable, while
--      removing @f@, because @f@ no longer allocates and can be floated to
--      top-level.
--   4. Actually float the binding of @f@ to top-level, eliminating the @let@
--      in the process.
--
-- This results in the following program (with free var annotations):
--
-- @
--    f x y a = a + x + y;
--    let x = ...; y = ... in
--    let g = {x y} \b -> f x y b + x in
--    g 5
-- @
--
-- This optimisation is all about lifting only when it is beneficial to do so.
-- The above seems like a worthwhile lift, judging from heap allocation:
-- We eliminate @f@'s closure, saving to allocate a closure with 2 words, while
-- not changing the size of @g@'s closure.
--
-- You can probably sense that there's some kind of cost model at play here.
-- And you are right! But we also employ a couple of other heuristics for the
-- lifting decision which are outlined in "StgLiftLams.Analysis#when".
--
-- The transformation is done in "StgLiftLams.Transformation", which calls out
-- to 'StgLiftLams.Analysis.goodToLift' for its lifting decision.
-- It relies on "StgLiftLams.LiftM", which abstracts some subtle STG invariants
-- into a monadic substrate.
--
-- Suffice to say: We trade heap allocation for stack allocation.
-- The additional arguments have to passed on the stack (or in registers,
-- depending on architecture) every time we call the function to save a single
-- heap allocation when entering the let binding. Nofib suggests a mean
-- improvement of about 1% for this pass, so it seems like a worthwhile thing to
-- do. Compile-times went up by 0.6%, so all in all a very modest change.
--
-- For a concrete example, look at @spectral/atom@. There's a call to 'zipWith'
-- that is ultimately compiled to something like this
-- (module desugaring/lowering to actual STG):
--
-- @
--    propagate dt = ...;
--    runExperiment ... =
--      let xs = ... in
--      let ys = ... in
--      let go = {dt go} \xs ys -> case (xs, ys) of
--            ([], []) -> []
--            (x:xs', y:ys') -> propagate dt x y : go xs' ys'
--      in go xs ys
-- @
--
-- This will lambda lift @go@ to top-level, speeding up the resulting program
-- by roughly one percent:
--
-- @
--    propagate dt = ...;
--    go dt xs ys = case (xs, ys) of
--      ([], []) -> []
--      (x:xs', y:ys') -> propagate dt x y : go dt xs' ys'
--    runExperiment ... =
--      let xs = ... in
--      let ys = ... in
--      in go dt xs ys
-- @


{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

\section[StgStats]{Gathers statistical information about programs}

The program gather statistics about
\begin{enumerate}
\item number of boxed cases
\item number of unboxed cases
\item number of let-no-escapes
\item number of non-updatable lets
\item number of updatable lets
\item number of applications
\item number of primitive applications
\item number of closures (does not include lets bound to constructors)
\item number of free variables in closures
%\item number of top-level functions
%\item number of top-level CAFs
\item number of constructors
\end{enumerate}
-}

{-# LANGUAGE CPP #-}

module StgStats ( showStgStats ) where

#include "HsVersions.h"

import GhcPrelude

import StgSyn

import Id (Id)
import Panic

import Data.Map (Map)
import qualified Data.Map as Map

data CounterType
  = Literals
  | Applications
  | ConstructorApps
  | PrimitiveApps
  | LetNoEscapes
  | StgCases
  | FreeVariables
  | ConstructorBinds Bool{-True<=>top-level-}
  | ReEntrantBinds   Bool{-ditto-}
  | SingleEntryBinds Bool{-ditto-}
  | UpdatableBinds   Bool{-ditto-}
  deriving (Eq, Ord)

type Count      = Int
type StatEnv    = Map CounterType Count

emptySE :: StatEnv
emptySE = Map.empty

combineSE :: StatEnv -> StatEnv -> StatEnv
combineSE = Map.unionWith (+)

combineSEs :: [StatEnv] -> StatEnv
combineSEs = foldr combineSE emptySE

countOne :: CounterType -> StatEnv
countOne c = Map.singleton c 1

{-
************************************************************************
*                                                                      *
\subsection{Top-level list of bindings (a ``program'')}
*                                                                      *
************************************************************************
-}

showStgStats :: [StgTopBinding] -> String

showStgStats prog
  = "STG Statistics:\n\n"
    ++ concat (map showc (Map.toList (gatherStgStats prog)))
  where
    showc (x,n) = (showString (s x) . shows n) "\n"

    s Literals                = "Literals                   "
    s Applications            = "Applications               "
    s ConstructorApps         = "ConstructorApps            "
    s PrimitiveApps           = "PrimitiveApps              "
    s LetNoEscapes            = "LetNoEscapes               "
    s StgCases                = "StgCases                   "
    s FreeVariables           = "FreeVariables              "
    s (ConstructorBinds True) = "ConstructorBinds_Top       "
    s (ReEntrantBinds True)   = "ReEntrantBinds_Top         "
    s (SingleEntryBinds True) = "SingleEntryBinds_Top       "
    s (UpdatableBinds True)   = "UpdatableBinds_Top         "
    s (ConstructorBinds _)    = "ConstructorBinds_Nested    "
    s (ReEntrantBinds _)      = "ReEntrantBindsBinds_Nested "
    s (SingleEntryBinds _)    = "SingleEntryBinds_Nested    "
    s (UpdatableBinds _)      = "UpdatableBinds_Nested      "

gatherStgStats :: [StgTopBinding] -> StatEnv
gatherStgStats binds = combineSEs (map statTopBinding binds)

{-
************************************************************************
*                                                                      *
\subsection{Bindings}
*                                                                      *
************************************************************************
-}

statTopBinding :: StgTopBinding -> StatEnv
statTopBinding (StgTopStringLit _ _) = countOne Literals
statTopBinding (StgTopLifted bind) = statBinding True bind

statBinding :: Bool -- True <=> top-level; False <=> nested
            -> StgBinding
            -> StatEnv

statBinding top (StgNonRec b rhs)
  = statRhs top (b, rhs)

statBinding top (StgRec pairs)
  = combineSEs (map (statRhs top) pairs)

statRhs :: Bool -> (Id, StgRhs) -> StatEnv

statRhs top (_, StgRhsCon _ _ _)
  = countOne (ConstructorBinds top)

statRhs top (_, StgRhsClosure _ _ u _ body)
  = statExpr body `combineSE`
    countOne (
      case u of
        ReEntrant   -> ReEntrantBinds   top
        Updatable   -> UpdatableBinds   top
        SingleEntry -> SingleEntryBinds top
    )

{-
************************************************************************
*                                                                      *
\subsection{Expressions}
*                                                                      *
************************************************************************
-}

statExpr :: StgExpr -> StatEnv

statExpr (StgApp _ _)     = countOne Applications
statExpr (StgLit _)       = countOne Literals
statExpr (StgConApp _ _ _)= countOne ConstructorApps
statExpr (StgOpApp _ _ _) = countOne PrimitiveApps
statExpr (StgTick _ e)    = statExpr e

statExpr (StgLetNoEscape _ binds body)
  = statBinding False{-not top-level-} binds    `combineSE`
    statExpr body                               `combineSE`
    countOne LetNoEscapes

statExpr (StgLet _ binds body)
  = statBinding False{-not top-level-} binds    `combineSE`
    statExpr body

statExpr (StgCase expr _ _ alts)
  = statExpr expr       `combineSE`
    stat_alts alts      `combineSE`
    countOne StgCases
  where
    stat_alts alts
        = combineSEs (map statExpr [ e | (_,_,e) <- alts ])

statExpr (StgLam {}) = panic "statExpr StgLam"


{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-2012

Note [Unarisation]
~~~~~~~~~~~~~~~~~~
The idea of this pass is to translate away *all* unboxed-tuple and unboxed-sum
binders. So for example:

  f (x :: (# Int, Bool #)) = f x + f (# 1, True #)

  ==>

  f (x1 :: Int) (x2 :: Bool) = f x1 x2 + f 1 True

It is important that we do this at the STG level and NOT at the Core level
because it would be very hard to make this pass Core-type-preserving. In this
example the type of 'f' changes, for example.

STG fed to the code generators *must* be unarised because the code generators do
not support unboxed tuple and unboxed sum binders natively.

In more detail: (see next note for unboxed sums)

Suppose that a variable x : (# t1, t2 #).

  * At the binding site for x, make up fresh vars  x1:t1, x2:t2

  * Extend the UnariseEnv   x :-> MultiVal [x1,x2]

  * Replace the binding with a curried binding for x1,x2

       Lambda:   \x.e                ==>   \x1 x2. e
       Case alt: MkT a b x c d -> e  ==>   MkT a b x1 x2 c d -> e

  * Replace argument occurrences with a sequence of args via a lookup in
    UnariseEnv

       f a b x c d   ==>   f a b x1 x2 c d

  * Replace tail-call occurrences with an unboxed tuple via a lookup in
    UnariseEnv

       x  ==>  (# x1, x2 #)

    So, for example

       f x = x    ==>   f x1 x2 = (# x1, x2 #)

  * We /always/ eliminate a case expression when

       - It scrutinises an unboxed tuple or unboxed sum

       - The scrutinee is a variable (or when it is an explicit tuple, but the
         simplifier eliminates those)

    The case alternative (there can be only one) can be one of these two
    things:

      - An unboxed tuple pattern. e.g.

          case v of x { (# x1, x2, x3 #) -> ... }

        Scrutinee has to be in form `(# t1, t2, t3 #)` so we just extend the
        environment with

          x :-> MultiVal [t1,t2,t3]
          x1 :-> UnaryVal t1, x2 :-> UnaryVal t2, x3 :-> UnaryVal t3

      - A DEFAULT alternative. Just the same, without the bindings for x1,x2,x3

By the end of this pass, we only have unboxed tuples in return positions.
Unboxed sums are completely eliminated, see next note.

Note [Translating unboxed sums to unboxed tuples]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Unarise also eliminates unboxed sum binders, and translates unboxed sums in
return positions to unboxed tuples. We want to overlap fields of a sum when
translating it to a tuple to have efficient memory layout. When translating a
sum pattern to a tuple pattern, we need to translate it so that binders of sum
alternatives will be mapped to right arguments after the term translation. So
translation of sum DataCon applications to tuple DataCon applications and
translation of sum patterns to tuple patterns need to be in sync.

These translations work like this. Suppose we have

  (# x1 | | ... #) :: (# t1 | t2 | ... #)

remember that t1, t2 ... can be sums and tuples too. So we first generate
layouts of those. Then we "merge" layouts of each alternative, which gives us a
sum layout with best overlapping possible.

Layout of a flat type 'ty1' is just [ty1].
Layout of a tuple is just concatenation of layouts of its fields.

For layout of a sum type,

  - We first get layouts of all alternatives.
  - We sort these layouts based on their "slot types".
  - We merge all the alternatives.

For example, say we have (# (# Int#, Char #) | (# Int#, Int# #) | Int# #)

  - Layouts of alternatives: [ [Word, Ptr], [Word, Word], [Word] ]
  - Sorted: [ [Ptr, Word], [Word, Word], [Word] ]
  - Merge all alternatives together: [ Ptr, Word, Word ]

We add a slot for the tag to the first position. So our tuple type is

  (# Tag#, Any, Word#, Word# #)
  (we use Any for pointer slots)

Now, any term of this sum type needs to generate a tuple of this type instead.
The translation works by simply putting arguments to first slots that they fit
in. Suppose we had

  (# (# 42#, 'c' #) | | #)

42# fits in Word#, 'c' fits in Any, so we generate this application:

  (# 1#, 'c', 42#, rubbish #)

Another example using the same type: (# | (# 2#, 3# #) | #). 2# fits in Word#,
3# fits in Word #, so we get:

  (# 2#, rubbish, 2#, 3# #).

Note [Types in StgConApp]
~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have this unboxed sum term:

  (# 123 | #)

What will be the unboxed tuple representation? We can't tell without knowing the
type of this term. For example, these are all valid tuples for this:

  (# 1#, 123 #)          -- when type is (# Int | String #)
  (# 1#, 123, rubbish #) -- when type is (# Int | Float# #)
  (# 1#, 123, rubbish, rubbish #)
                         -- when type is (# Int | (# Int, Int, Int #) #)

So we pass type arguments of the DataCon's TyCon in StgConApp to decide what
layout to use. Note that unlifted values can't be let-bound, so we don't need
types in StgRhsCon.

Note [UnariseEnv can map to literals]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
To avoid redundant case expressions when unarising unboxed sums, UnariseEnv
needs to map variables to literals too. Suppose we have this Core:

  f (# x | #)

  ==> (CorePrep)

  case (# x | #) of y {
    _ -> f y
  }

  ==> (MultiVal)

  case (# 1#, x #) of [x1, x2] {
    _ -> f x1 x2
  }

To eliminate this case expression we need to map x1 to 1# in UnariseEnv:

  x1 :-> UnaryVal 1#, x2 :-> UnaryVal x

so that `f x1 x2` becomes `f 1# x`.

Note [Unarisation and arity]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Because of unarisation, the arity that will be recorded in the generated info
table for an Id may be larger than the idArity. Instead we record what we call
the RepArity, which is the Arity taking into account any expanded arguments, and
corresponds to the number of (possibly-void) *registers* arguments will arrive
in.

Note [Post-unarisation invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
STG programs after unarisation have these invariants:

  * No unboxed sums at all.

  * No unboxed tuple binders. Tuples only appear in return position.

  * DataCon applications (StgRhsCon and StgConApp) don't have void arguments.
    This means that it's safe to wrap `StgArg`s of DataCon applications with
    `StgCmmEnv.NonVoid`, for example.

  * Alt binders (binders in patterns) are always non-void.
-}

{-# LANGUAGE CPP, TupleSections #-}

module UnariseStg (unarise) where

#include "HsVersions.h"

import GhcPrelude

import BasicTypes
import CoreSyn
import DataCon
import FastString (FastString, mkFastString)
import Id
import Literal
import MkCore (aBSENT_SUM_FIELD_ERROR_ID)
import MkId (voidPrimId, voidArgId)
import MonadUtils (mapAccumLM)
import Outputable
import RepType
import StgSyn
import Type
import TysPrim (intPrimTy,wordPrimTy,word64PrimTy)
import TysWiredIn
import UniqSupply
import Util
import VarEnv

import Data.Bifunctor (second)
import Data.Maybe (mapMaybe)
import qualified Data.IntMap as IM

--------------------------------------------------------------------------------

-- | A mapping from binders to the Ids they were expanded/renamed to.
--
--   x :-> MultiVal [a,b,c] in rho
--
-- iff  x's typePrimRep is not a singleton, or equivalently
--      x's type is an unboxed tuple, sum or void.
--
--    x :-> UnaryVal x'
--
-- iff x's RepType is UnaryRep or equivalently
--     x's type is not unboxed tuple, sum or void.
--
-- So
--     x :-> MultiVal [a] in rho
-- means x is represented by singleton tuple.
--
--     x :-> MultiVal [] in rho
-- means x is void.
--
-- INVARIANT: OutStgArgs in the range only have NvUnaryTypes
--            (i.e. no unboxed tuples, sums or voids)
--
type UnariseEnv = VarEnv UnariseVal

data UnariseVal
  = MultiVal [OutStgArg] -- MultiVal to tuple. Can be empty list (void).
  | UnaryVal OutStgArg   -- See NOTE [Renaming during unarisation].

instance Outputable UnariseVal where
  ppr (MultiVal args) = text "MultiVal" <+> ppr args
  ppr (UnaryVal arg)   = text "UnaryVal" <+> ppr arg

-- | Extend the environment, checking the UnariseEnv invariant.
extendRho :: UnariseEnv -> Id -> UnariseVal -> UnariseEnv
extendRho rho x (MultiVal args)
  = ASSERT(all (isNvUnaryType . stgArgType) args)
    extendVarEnv rho x (MultiVal args)
extendRho rho x (UnaryVal val)
  = ASSERT(isNvUnaryType (stgArgType val))
    extendVarEnv rho x (UnaryVal val)

--------------------------------------------------------------------------------

unarise :: UniqSupply -> [StgTopBinding] -> [StgTopBinding]
unarise us binds = initUs_ us (mapM (unariseTopBinding emptyVarEnv) binds)

unariseTopBinding :: UnariseEnv -> StgTopBinding -> UniqSM StgTopBinding
unariseTopBinding rho (StgTopLifted bind)
  = StgTopLifted <$> unariseBinding rho bind
unariseTopBinding _ bind@StgTopStringLit{} = return bind

unariseBinding :: UnariseEnv -> StgBinding -> UniqSM StgBinding
unariseBinding rho (StgNonRec x rhs)
  = StgNonRec x <$> unariseRhs rho rhs
unariseBinding rho (StgRec xrhss)
  = StgRec <$> mapM (\(x, rhs) -> (x,) <$> unariseRhs rho rhs) xrhss

unariseRhs :: UnariseEnv -> StgRhs -> UniqSM StgRhs
unariseRhs rho (StgRhsClosure ext ccs update_flag args expr)
  = do (rho', args1) <- unariseFunArgBinders rho args
       expr' <- unariseExpr rho' expr
       return (StgRhsClosure ext ccs update_flag args1 expr')

unariseRhs rho (StgRhsCon ccs con args)
  = ASSERT(not (isUnboxedTupleCon con || isUnboxedSumCon con))
    return (StgRhsCon ccs con (unariseConArgs rho args))

--------------------------------------------------------------------------------

unariseExpr :: UnariseEnv -> StgExpr -> UniqSM StgExpr

unariseExpr rho e@(StgApp f [])
  = case lookupVarEnv rho f of
      Just (MultiVal args)  -- Including empty tuples
        -> return (mkTuple args)
      Just (UnaryVal (StgVarArg f'))
        -> return (StgApp f' [])
      Just (UnaryVal (StgLitArg f'))
        -> return (StgLit f')
      Nothing
        -> return e

unariseExpr rho e@(StgApp f args)
  = return (StgApp f' (unariseFunArgs rho args))
  where
    f' = case lookupVarEnv rho f of
           Just (UnaryVal (StgVarArg f')) -> f'
           Nothing -> f
           err -> pprPanic "unariseExpr - app2" (ppr e $$ ppr err)
               -- Can't happen because 'args' is non-empty, and
               -- a tuple or sum cannot be applied to anything

unariseExpr _ (StgLit l)
  = return (StgLit l)

unariseExpr rho (StgConApp dc args ty_args)
  | Just args' <- unariseMulti_maybe rho dc args ty_args
  = return (mkTuple args')

  | otherwise
  , let args' = unariseConArgs rho args
  = return (StgConApp dc args' (map stgArgType args'))

unariseExpr rho (StgOpApp op args ty)
  = return (StgOpApp op (unariseFunArgs rho args) ty)

unariseExpr _ e@StgLam{}
  = pprPanic "unariseExpr: found lambda" (ppr e)

unariseExpr rho (StgCase scrut bndr alt_ty alts)
  -- tuple/sum binders in the scrutinee can always be eliminated
  | StgApp v [] <- scrut
  , Just (MultiVal xs) <- lookupVarEnv rho v
  = elimCase rho xs bndr alt_ty alts

  -- Handle strict lets for tuples and sums:
  --   case (# a,b #) of r -> rhs
  -- and analogously for sums
  | StgConApp dc args ty_args <- scrut
  , Just args' <- unariseMulti_maybe rho dc args ty_args
  = elimCase rho args' bndr alt_ty alts

  -- general case
  | otherwise
  = do scrut' <- unariseExpr rho scrut
       alts'  <- unariseAlts rho alt_ty bndr alts
       return (StgCase scrut' bndr alt_ty alts')
                       -- bndr may have a unboxed sum/tuple type but it will be
                       -- dead after unarise (checked in StgLint)

unariseExpr rho (StgLet ext bind e)
  = StgLet ext <$> unariseBinding rho bind <*> unariseExpr rho e

unariseExpr rho (StgLetNoEscape ext bind e)
  = StgLetNoEscape ext <$> unariseBinding rho bind <*> unariseExpr rho e

unariseExpr rho (StgTick tick e)
  = StgTick tick <$> unariseExpr rho e

-- Doesn't return void args.
unariseMulti_maybe :: UnariseEnv -> DataCon -> [InStgArg] -> [Type] -> Maybe [OutStgArg]
unariseMulti_maybe rho dc args ty_args
  | isUnboxedTupleCon dc
  = Just (unariseConArgs rho args)

  | isUnboxedSumCon dc
  , let args1 = ASSERT(isSingleton args) (unariseConArgs rho args)
  = Just (mkUbxSum dc ty_args args1)

  | otherwise
  = Nothing

--------------------------------------------------------------------------------

elimCase :: UnariseEnv
         -> [OutStgArg] -- non-void args
         -> InId -> AltType -> [InStgAlt] -> UniqSM OutStgExpr

elimCase rho args bndr (MultiValAlt _) [(_, bndrs, rhs)]
  = do let rho1 = extendRho rho bndr (MultiVal args)
           rho2
             | isUnboxedTupleBndr bndr
             = mapTupleIdBinders bndrs args rho1
             | otherwise
             = ASSERT(isUnboxedSumBndr bndr)
               if null bndrs then rho1
                             else mapSumIdBinders bndrs args rho1

       unariseExpr rho2 rhs

elimCase rho args bndr (MultiValAlt _) alts
  | isUnboxedSumBndr bndr
  = do let (tag_arg : real_args) = args
       tag_bndr <- mkId (mkFastString "tag") tagTy
          -- this won't be used but we need a binder anyway
       let rho1 = extendRho rho bndr (MultiVal args)
           scrut' = case tag_arg of
                      StgVarArg v     -> StgApp v []
                      StgLitArg l     -> StgLit l

       alts' <- unariseSumAlts rho1 real_args alts
       return (StgCase scrut' tag_bndr tagAltTy alts')

elimCase _ args bndr alt_ty alts
  = pprPanic "elimCase - unhandled case"
      (ppr args <+> ppr bndr <+> ppr alt_ty $$ ppr alts)

--------------------------------------------------------------------------------

unariseAlts :: UnariseEnv -> AltType -> InId -> [StgAlt] -> UniqSM [StgAlt]
unariseAlts rho (MultiValAlt n) bndr [(DEFAULT, [], e)]
  | isUnboxedTupleBndr bndr
  = do (rho', ys) <- unariseConArgBinder rho bndr
       e' <- unariseExpr rho' e
       return [(DataAlt (tupleDataCon Unboxed n), ys, e')]

unariseAlts rho (MultiValAlt n) bndr [(DataAlt _, ys, e)]
  | isUnboxedTupleBndr bndr
  = do (rho', ys1) <- unariseConArgBinders rho ys
       MASSERT(ys1 `lengthIs` n)
       let rho'' = extendRho rho' bndr (MultiVal (map StgVarArg ys1))
       e' <- unariseExpr rho'' e
       return [(DataAlt (tupleDataCon Unboxed n), ys1, e')]

unariseAlts _ (MultiValAlt _) bndr alts
  | isUnboxedTupleBndr bndr
  = pprPanic "unariseExpr: strange multi val alts" (ppr alts)

-- In this case we don't need to scrutinize the tag bit
unariseAlts rho (MultiValAlt _) bndr [(DEFAULT, _, rhs)]
  | isUnboxedSumBndr bndr
  = do (rho_sum_bndrs, sum_bndrs) <- unariseConArgBinder rho bndr
       rhs' <- unariseExpr rho_sum_bndrs rhs
       return [(DataAlt (tupleDataCon Unboxed (length sum_bndrs)), sum_bndrs, rhs')]

unariseAlts rho (MultiValAlt _) bndr alts
  | isUnboxedSumBndr bndr
  = do (rho_sum_bndrs, scrt_bndrs@(tag_bndr : real_bndrs)) <- unariseConArgBinder rho bndr
       alts' <- unariseSumAlts rho_sum_bndrs (map StgVarArg real_bndrs) alts
       let inner_case = StgCase (StgApp tag_bndr []) tag_bndr tagAltTy alts'
       return [ (DataAlt (tupleDataCon Unboxed (length scrt_bndrs)),
                 scrt_bndrs,
                 inner_case) ]

unariseAlts rho _ _ alts
  = mapM (\alt -> unariseAlt rho alt) alts

unariseAlt :: UnariseEnv -> StgAlt -> UniqSM StgAlt
unariseAlt rho (con, xs, e)
  = do (rho', xs') <- unariseConArgBinders rho xs
       (con, xs',) <$> unariseExpr rho' e

--------------------------------------------------------------------------------

-- | Make alternatives that match on the tag of a sum
-- (i.e. generate LitAlts for the tag)
unariseSumAlts :: UnariseEnv
               -> [StgArg] -- sum components _excluding_ the tag bit.
               -> [StgAlt] -- original alternative with sum LHS
               -> UniqSM [StgAlt]
unariseSumAlts env args alts
  = do alts' <- mapM (unariseSumAlt env args) alts
       return (mkDefaultLitAlt alts')

unariseSumAlt :: UnariseEnv
              -> [StgArg] -- sum components _excluding_ the tag bit.
              -> StgAlt   -- original alternative with sum LHS
              -> UniqSM StgAlt
unariseSumAlt rho _ (DEFAULT, _, e)
  = ( DEFAULT, [], ) <$> unariseExpr rho e

unariseSumAlt rho args (DataAlt sumCon, bs, e)
  = do let rho' = mapSumIdBinders bs args rho
       e' <- unariseExpr rho' e
       return ( LitAlt (LitNumber LitNumInt (fromIntegral (dataConTag sumCon)) intPrimTy), [], e' )

unariseSumAlt _ scrt alt
  = pprPanic "unariseSumAlt" (ppr scrt $$ ppr alt)

--------------------------------------------------------------------------------

mapTupleIdBinders
  :: [InId]       -- Un-processed binders of a tuple alternative.
                  -- Can have void binders.
  -> [OutStgArg]  -- Arguments that form the tuple (after unarisation).
                  -- Can't have void args.
  -> UnariseEnv
  -> UnariseEnv
mapTupleIdBinders ids args0 rho0
  = ASSERT(not (any (isVoidTy . stgArgType) args0))
    let
      ids_unarised :: [(Id, [PrimRep])]
      ids_unarised = map (\id -> (id, typePrimRep (idType id))) ids

      map_ids :: UnariseEnv -> [(Id, [PrimRep])] -> [StgArg] -> UnariseEnv
      map_ids rho [] _  = rho
      map_ids rho ((x, x_reps) : xs) args =
        let
          x_arity = length x_reps
          (x_args, args') =
            ASSERT(args `lengthAtLeast` x_arity)
            splitAt x_arity args

          rho'
            | x_arity == 1
            = ASSERT(x_args `lengthIs` 1)
              extendRho rho x (UnaryVal (head x_args))
            | otherwise
            = extendRho rho x (MultiVal x_args)
        in
          map_ids rho' xs args'
    in
      map_ids rho0 ids_unarised args0

mapSumIdBinders
  :: [InId]      -- Binder of a sum alternative (remember that sum patterns
                 -- only have one binder, so this list should be a singleton)
  -> [OutStgArg] -- Arguments that form the sum (NOT including the tag).
                 -- Can't have void args.
  -> UnariseEnv
  -> UnariseEnv

mapSumIdBinders [id] args rho0
  = ASSERT(not (any (isVoidTy . stgArgType) args))
    let
      arg_slots = map primRepSlot $ concatMap (typePrimRep . stgArgType) args
      id_slots  = map primRepSlot $ typePrimRep (idType id)
      layout1   = layoutUbxSum arg_slots id_slots
    in
      if isMultiValBndr id
        then extendRho rho0 id (MultiVal [ args !! i | i <- layout1 ])
        else ASSERT(layout1 `lengthIs` 1)
             extendRho rho0 id (UnaryVal (args !! head layout1))

mapSumIdBinders ids sum_args _
  = pprPanic "mapSumIdBinders" (ppr ids $$ ppr sum_args)

-- | Build a unboxed sum term from arguments of an alternative.
--
-- Example, for (# x | #) :: (# (# #) | Int #) we call
--
--   mkUbxSum (# _ | #) [ (# #), Int ] [ voidPrimId ]
--
-- which returns
--
--   [ 1#, rubbish ]
--
mkUbxSum
  :: DataCon      -- Sum data con
  -> [Type]       -- Type arguments of the sum data con
  -> [OutStgArg]  -- Actual arguments of the alternative.
  -> [OutStgArg]  -- Final tuple arguments
mkUbxSum dc ty_args args0
  = let
      (_ : sum_slots) = ubxSumRepType (map typePrimRep ty_args)
        -- drop tag slot

      tag = dataConTag dc

      layout'  = layoutUbxSum sum_slots (mapMaybe (typeSlotTy . stgArgType) args0)
      tag_arg  = StgLitArg (LitNumber LitNumInt (fromIntegral tag) intPrimTy)
      arg_idxs = IM.fromList (zipEqual "mkUbxSum" layout' args0)

      mkTupArgs :: Int -> [SlotTy] -> IM.IntMap StgArg -> [StgArg]
      mkTupArgs _ [] _
        = []
      mkTupArgs arg_idx (slot : slots_left) arg_map
        | Just stg_arg <- IM.lookup arg_idx arg_map
        = stg_arg : mkTupArgs (arg_idx + 1) slots_left arg_map
        | otherwise
        = slotRubbishArg slot : mkTupArgs (arg_idx + 1) slots_left arg_map

      slotRubbishArg :: SlotTy -> StgArg
      slotRubbishArg PtrSlot    = StgVarArg aBSENT_SUM_FIELD_ERROR_ID
                         -- See Note [aBSENT_SUM_FIELD_ERROR_ID] in MkCore
      slotRubbishArg WordSlot   = StgLitArg (LitNumber LitNumWord 0 wordPrimTy)
      slotRubbishArg Word64Slot = StgLitArg (LitNumber LitNumWord64 0 word64PrimTy)
      slotRubbishArg FloatSlot  = StgLitArg (LitFloat 0)
      slotRubbishArg DoubleSlot = StgLitArg (LitDouble 0)
    in
      tag_arg : mkTupArgs 0 sum_slots arg_idxs

--------------------------------------------------------------------------------

{-
For arguments (StgArg) and binders (Id) we have two kind of unarisation:

  - When unarising function arg binders and arguments, we don't want to remove
    void binders and arguments. For example,

      f :: (# (# #), (# #) #) -> Void# -> RealWorld# -> ...
      f x y z = <body>

    Here after unarise we should still get a function with arity 3. Similarly
    in the call site we shouldn't remove void arguments:

      f (# (# #), (# #) #) voidId rw

    When unarising <body>, we extend the environment with these binders:

      x :-> MultiVal [], y :-> MultiVal [], z :-> MultiVal []

    Because their rep types are `MultiRep []` (aka. void). This means that when
    we see `x` in a function argument position, we actually replace it with a
    void argument. When we see it in a DataCon argument position, we just get
    rid of it, because DataCon applications in STG are always saturated.

  - When unarising case alternative binders we remove void binders, but we
    still update the environment the same way, because those binders may be
    used in the RHS. Example:

      case x of y {
        (# x1, x2, x3 #) -> <RHS>
      }

    We know that y can't be void, because we don't scrutinize voids, so x will
    be unarised to some number of arguments, and those arguments will have at
    least one non-void thing. So in the rho we will have something like:

      x :-> MultiVal [xu1, xu2]

    Now, after we eliminate void binders in the pattern, we get exactly the same
    number of binders, and extend rho again with these:

      x1 :-> UnaryVal xu1
      x2 :-> MultiVal [] -- x2 is void
      x3 :-> UnaryVal xu2

    Now when we see x2 in a function argument position or in return position, we
    generate void#. In constructor argument position, we just remove it.

So in short, when we have a void id,

  - We keep it if it's a lambda argument binder or
                       in argument position of an application.

  - We remove it if it's a DataCon field binder or
                         in argument position of a DataCon application.
-}

unariseArgBinder
    :: Bool -- data con arg?
    -> UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
unariseArgBinder is_con_arg rho x =
  case typePrimRep (idType x) of
    []
      | is_con_arg
      -> return (extendRho rho x (MultiVal []), [])
      | otherwise -- fun arg, do not remove void binders
      -> return (extendRho rho x (MultiVal []), [voidArgId])

    [rep]
      -- Arg represented as single variable, but original type may still be an
      -- unboxed sum/tuple, e.g. (# Void# | Void# #).
      --
      -- While not unarising the binder in this case does not break any programs
      -- (because it unarises to a single variable), it triggers StgLint as we
      -- break the the post-unarisation invariant that says unboxed tuple/sum
      -- binders should vanish. See Note [Post-unarisation invariants].
      | isUnboxedSumType (idType x) || isUnboxedTupleType (idType x)
      -> do x' <- mkId (mkFastString "us") (primRepToType rep)
            return (extendRho rho x (MultiVal [StgVarArg x']), [x'])
      | otherwise
      -> return (rho, [x])

    reps -> do
      xs <- mkIds (mkFastString "us") (map primRepToType reps)
      return (extendRho rho x (MultiVal (map StgVarArg xs)), xs)

--------------------------------------------------------------------------------

-- | MultiVal a function argument. Never returns an empty list.
unariseFunArg :: UnariseEnv -> StgArg -> [StgArg]
unariseFunArg rho (StgVarArg x) =
  case lookupVarEnv rho x of
    Just (MultiVal [])  -> [voidArg]   -- NB: do not remove void args
    Just (MultiVal as)  -> as
    Just (UnaryVal arg) -> [arg]
    Nothing             -> [StgVarArg x]
unariseFunArg _ arg = [arg]

unariseFunArgs :: UnariseEnv -> [StgArg] -> [StgArg]
unariseFunArgs = concatMap . unariseFunArg

unariseFunArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
unariseFunArgBinders rho xs = second concat <$> mapAccumLM unariseFunArgBinder rho xs

-- Result list of binders is never empty
unariseFunArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
unariseFunArgBinder = unariseArgBinder False

--------------------------------------------------------------------------------

-- | MultiVal a DataCon argument. Returns an empty list when argument is void.
unariseConArg :: UnariseEnv -> InStgArg -> [OutStgArg]
unariseConArg rho (StgVarArg x) =
  case lookupVarEnv rho x of
    Just (UnaryVal arg) -> [arg]
    Just (MultiVal as) -> as      -- 'as' can be empty
    Nothing
      | isVoidTy (idType x) -> [] -- e.g. C realWorld#
                                  -- Here realWorld# is not in the envt, but
                                  -- is a void, and so should be eliminated
      | otherwise -> [StgVarArg x]
unariseConArg _ arg@(StgLitArg lit) =
    ASSERT(not (isVoidTy (literalType lit)))  -- We have no void literals
    [arg]

unariseConArgs :: UnariseEnv -> [InStgArg] -> [OutStgArg]
unariseConArgs = concatMap . unariseConArg

unariseConArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
unariseConArgBinders rho xs = second concat <$> mapAccumLM unariseConArgBinder rho xs

-- Different from `unariseFunArgBinder`: result list of binders may be empty.
-- See DataCon applications case in Note [Post-unarisation invariants].
unariseConArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
unariseConArgBinder = unariseArgBinder True

--------------------------------------------------------------------------------

mkIds :: FastString -> [UnaryType] -> UniqSM [Id]
mkIds fs tys = mapM (mkId fs) tys

mkId :: FastString -> UnaryType -> UniqSM Id
mkId = mkSysLocalOrCoVarM

isMultiValBndr :: Id -> Bool
isMultiValBndr id
  | [_] <- typePrimRep (idType id)
  = False
  | otherwise
  = True

isUnboxedSumBndr :: Id -> Bool
isUnboxedSumBndr = isUnboxedSumType . idType

isUnboxedTupleBndr :: Id -> Bool
isUnboxedTupleBndr = isUnboxedTupleType . idType

mkTuple :: [StgArg] -> StgExpr
mkTuple args = StgConApp (tupleDataCon Unboxed (length args)) args (map stgArgType args)

tagAltTy :: AltType
tagAltTy = PrimAlt IntRep

tagTy :: Type
tagTy = intPrimTy

voidArg :: StgArg
voidArg = StgVarArg voidPrimId

mkDefaultLitAlt :: [StgAlt] -> [StgAlt]
-- We have an exhauseive list of literal alternatives
--    1# -> e1
--    2# -> e2
-- Since they are exhaustive, we can replace one with DEFAULT, to avoid
-- generating a final test. Remember, the DEFAULT comes first if it exists.
mkDefaultLitAlt [] = pprPanic "elimUbxSumExpr.mkDefaultAlt" (text "Empty alts")
mkDefaultLitAlt alts@((DEFAULT, _, _) : _) = alts
mkDefaultLitAlt ((LitAlt{}, [], rhs) : alts) = (DEFAULT, [], rhs) : alts
mkDefaultLitAlt alts = pprPanic "mkDefaultLitAlt" (text "Not a lit alt:" <+> ppr alts)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds #-}

-- | Provides the heuristics for when it's beneficial to lambda lift bindings.
-- Most significantly, this employs a cost model to estimate impact on heap
-- allocations, by looking at an STG expression's 'Skeleton'.
module StgLiftLams.Analysis (
    -- * #when# When to lift
    -- $when

    -- * #clogro# Estimating closure growth
    -- $clogro

    -- * AST annotation
    Skeleton(..), BinderInfo(..), binderInfoBndr,
    LlStgBinding, LlStgExpr, LlStgRhs, LlStgAlt, tagSkeletonTopBind,
    -- * Lifting decision
    goodToLift,
    closureGrowth -- Exported just for the docs
  ) where

import GhcPrelude

import BasicTypes
import Demand
import DynFlags
import Id
import SMRep ( WordOff )
import StgSyn
import qualified StgCmmArgRep
import qualified StgCmmClosure
import qualified StgCmmLayout
import Outputable
import Util
import VarSet

import Data.Maybe ( mapMaybe )

-- Note [When to lift]
-- ~~~~~~~~~~~~~~~~~~~
-- $when
-- The analysis proceeds in two steps:
--
--   1. It tags the syntax tree with analysis information in the form of
--      'BinderInfo' at each binder and 'Skeleton's at each let-binding
--      by 'tagSkeletonTopBind' and friends.
--   2. The resulting syntax tree is treated by the "StgLiftLams.Transformation"
--      module, calling out to 'goodToLift' to decide if a binding is worthwhile
--      to lift.
--      'goodToLift' consults argument occurrence information in 'BinderInfo'
--      and estimates 'closureGrowth', for which it needs the 'Skeleton'.
--
-- So the annotations from 'tagSkeletonTopBind' ultimately fuel 'goodToLift',
-- which employs a number of heuristics to identify and exclude lambda lifting
-- opportunities deemed non-beneficial:
--
--  [Top-level bindings] can't be lifted.
--  [Thunks] and data constructors shouldn't be lifted in order not to destroy
--    sharing.
--  [Argument occurrences] #arg_occs# of binders prohibit them to be lifted.
--    Doing the lift would re-introduce the very allocation at call sites that
--    we tried to get rid off in the first place. We capture analysis
--    information in 'BinderInfo'. Note that we also consider a nullary
--    application as argument occurrence, because it would turn into an n-ary
--    partial application created by a generic apply function. This occurs in
--    CPS-heavy code like the CS benchmark.
--  [Join points] should not be lifted, simply because there's no reduction in
--    allocation to be had.
--  [Abstracting over join points] destroys join points, because they end up as
--    arguments to the lifted function.
--  [Abstracting over known local functions] turns a known call into an unknown
--    call (e.g. some @stg_ap_*@), which is generally slower. Can be turned off
--    with @-fstg-lift-lams-known@.
--  [Calling convention] Don't lift when the resulting function would have a
--    higher arity than available argument registers for the calling convention.
--    Can be influenced with @-fstg-lift-(non)rec-args(-any)@.
--  [Closure growth] introduced when former free variables have to be available
--    at call sites may actually lead to an increase in overall allocations
--  resulting from a lift. Estimating closure growth is described in
--  "StgLiftLams.Analysis#clogro" and is what most of this module is ultimately
--  concerned with.
--
-- There's a <https://ghc.haskell.org/trac/ghc/wiki/LateLamLift wiki page> with
-- some more background and history.

-- Note [Estimating closure growth]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- $clogro
-- We estimate closure growth by abstracting the syntax tree into a 'Skeleton',
-- capturing only syntactic details relevant to 'closureGrowth', such as
--
--   * 'ClosureSk', representing closure allocation.
--   * 'RhsSk', representing a RHS of a binding and how many times it's called
--     by an appropriate 'DmdShell'.
--   * 'AltSk', 'BothSk' and 'NilSk' for choice, sequence and empty element.
--
-- This abstraction is mostly so that the main analysis function 'closureGrowth'
-- can stay simple and focused. Also, skeletons tend to be much smaller than
-- the syntax tree they abstract, so it makes sense to construct them once and
-- and operate on them instead of the actual syntax tree.
--
-- A more detailed treatment of computing closure growth, including examples,
-- can be found in the paper referenced from the
-- <https://ghc.haskell.org/trac/ghc/wiki/LateLamLift wiki page>.

llTrace :: String -> SDoc -> a -> a
llTrace _ _ c = c
-- llTrace a b c = pprTrace a b c

type instance BinderP      'LiftLams = BinderInfo
type instance XRhsClosure  'LiftLams = DIdSet
type instance XLet         'LiftLams = Skeleton
type instance XLetNoEscape 'LiftLams = Skeleton

freeVarsOfRhs :: (XRhsClosure pass ~ DIdSet) => GenStgRhs pass -> DIdSet
freeVarsOfRhs (StgRhsCon _ _ args) = mkDVarSet [ id | StgVarArg id <- args ]
freeVarsOfRhs (StgRhsClosure fvs _ _ _ _) = fvs

-- | Captures details of the syntax tree relevant to the cost model, such as
-- closures, multi-shot lambdas and case expressions.
data Skeleton
  = ClosureSk !Id !DIdSet {- ^ free vars -} !Skeleton
  | RhsSk !DmdShell {- ^ how often the RHS was entered -} !Skeleton
  | AltSk !Skeleton !Skeleton
  | BothSk !Skeleton !Skeleton
  | NilSk

bothSk :: Skeleton -> Skeleton -> Skeleton
bothSk NilSk b = b
bothSk a NilSk = a
bothSk a b     = BothSk a b

altSk :: Skeleton -> Skeleton -> Skeleton
altSk NilSk b = b
altSk a NilSk = a
altSk a b     = AltSk a b

rhsSk :: DmdShell -> Skeleton -> Skeleton
rhsSk _        NilSk = NilSk
rhsSk body_dmd skel  = RhsSk body_dmd skel

-- | The type used in binder positions in 'GenStgExpr's.
data BinderInfo
  = BindsClosure !Id !Bool -- ^ Let(-no-escape)-bound thing with a flag
                           --   indicating whether it occurs as an argument
                           --   or in a nullary application
                           --   (see "StgLiftLams.Analysis#arg_occs").
  | BoringBinder !Id       -- ^ Every other kind of binder

-- | Gets the bound 'Id' out a 'BinderInfo'.
binderInfoBndr :: BinderInfo -> Id
binderInfoBndr (BoringBinder bndr)   = bndr
binderInfoBndr (BindsClosure bndr _) = bndr

-- | Returns 'Nothing' for 'BoringBinder's and 'Just' the flag indicating
-- occurrences as argument or in a nullary applications otherwise.
binderInfoOccursAsArg :: BinderInfo -> Maybe Bool
binderInfoOccursAsArg BoringBinder{}     = Nothing
binderInfoOccursAsArg (BindsClosure _ b) = Just b

instance Outputable Skeleton where
  ppr NilSk = text ""
  ppr (AltSk l r) = vcat
    [ text "{ " <+> ppr l
    , text "ALT"
    , text "  " <+> ppr r
    , text "}"
    ]
  ppr (BothSk l r) = ppr l $$ ppr r
  ppr (ClosureSk f fvs body) = ppr f <+> ppr fvs $$ nest 2 (ppr body)
  ppr (RhsSk body_dmd body) = hcat
    [ text "Î»["
    , ppr str
    , text ", "
    , ppr use
    , text "]. "
    , ppr body
    ]
    where
      str
        | isStrictDmd body_dmd = '1'
        | otherwise = '0'
      use
        | isAbsDmd body_dmd = '0'
        | isUsedOnce body_dmd = '1'
        | otherwise = 'Ï‰'

instance Outputable BinderInfo where
  ppr = ppr . binderInfoBndr

instance OutputableBndr BinderInfo where
  pprBndr b = pprBndr b . binderInfoBndr
  pprPrefixOcc = pprPrefixOcc . binderInfoBndr
  pprInfixOcc = pprInfixOcc . binderInfoBndr
  bndrIsJoin_maybe = bndrIsJoin_maybe . binderInfoBndr

mkArgOccs :: [StgArg] -> IdSet
mkArgOccs = mkVarSet . mapMaybe stg_arg_var
  where
    stg_arg_var (StgVarArg occ) = Just occ
    stg_arg_var _               = Nothing

-- | Tags every binder with its 'BinderInfo' and let bindings with their
-- 'Skeleton's.
tagSkeletonTopBind :: CgStgBinding -> LlStgBinding
-- NilSk is OK when tagging top-level bindings. Also, top-level things are never
-- lambda-lifted, so no need to track their argument occurrences. They can also
-- never be let-no-escapes (thus we pass False).
tagSkeletonTopBind bind = bind'
  where
    (_, _, _, bind') = tagSkeletonBinding False NilSk emptyVarSet bind

-- | Tags binders of an 'StgExpr' with its 'BinderInfo' and let bindings with
-- their 'Skeleton's. Additionally, returns its 'Skeleton' and the set of binder
-- occurrences in argument and nullary application position
-- (cf. "StgLiftLams.Analysis#arg_occs").
tagSkeletonExpr :: CgStgExpr -> (Skeleton, IdSet, LlStgExpr)
tagSkeletonExpr (StgLit lit)
  = (NilSk, emptyVarSet, StgLit lit)
tagSkeletonExpr (StgConApp con args tys)
  = (NilSk, mkArgOccs args, StgConApp con args tys)
tagSkeletonExpr (StgOpApp op args ty)
  = (NilSk, mkArgOccs args, StgOpApp op args ty)
tagSkeletonExpr (StgApp f args)
  = (NilSk, arg_occs, StgApp f args)
  where
    arg_occs
      -- This checks for nullary applications, which we treat the same as
      -- argument occurrences, see "StgLiftLams.Analysis#arg_occs".
      | null args = unitVarSet f
      | otherwise = mkArgOccs args
tagSkeletonExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam")
tagSkeletonExpr (StgCase scrut bndr ty alts)
  = (skel, arg_occs, StgCase scrut' bndr' ty alts')
  where
    (scrut_skel, scrut_arg_occs, scrut') = tagSkeletonExpr scrut
    (alt_skels, alt_arg_occss, alts') = mapAndUnzip3 tagSkeletonAlt alts
    skel = bothSk scrut_skel (foldr altSk NilSk alt_skels)
    arg_occs = unionVarSets (scrut_arg_occs:alt_arg_occss) `delVarSet` bndr
    bndr' = BoringBinder bndr
tagSkeletonExpr (StgTick t e)
  = (skel, arg_occs, StgTick t e')
  where
    (skel, arg_occs, e') = tagSkeletonExpr e
tagSkeletonExpr (StgLet _ bind body) = tagSkeletonLet False body bind
tagSkeletonExpr (StgLetNoEscape _ bind body) = tagSkeletonLet True body bind

mkLet :: Bool -> Skeleton -> LlStgBinding -> LlStgExpr -> LlStgExpr
mkLet True = StgLetNoEscape
mkLet _    = StgLet

tagSkeletonLet
  :: Bool
  -- ^ Is the binding a let-no-escape?
  -> CgStgExpr
  -- ^ Let body
  -> CgStgBinding
  -- ^ Binding group
  -> (Skeleton, IdSet, LlStgExpr)
  -- ^ RHS skeletons, argument occurrences and annotated binding
tagSkeletonLet is_lne body bind
  = (let_skel, arg_occs, mkLet is_lne scope bind' body')
  where
    (body_skel, body_arg_occs, body') = tagSkeletonExpr body
    (let_skel, arg_occs, scope, bind')
      = tagSkeletonBinding is_lne body_skel body_arg_occs bind

tagSkeletonBinding
  :: Bool
  -- ^ Is the binding a let-no-escape?
  -> Skeleton
  -- ^ Let body skeleton
  -> IdSet
  -- ^ Argument occurrences in the body
  -> CgStgBinding
  -- ^ Binding group
  -> (Skeleton, IdSet, Skeleton, LlStgBinding)
  -- ^ Let skeleton, argument occurrences, scope skeleton of binding and
  --   the annotated binding
tagSkeletonBinding is_lne body_skel body_arg_occs (StgNonRec bndr rhs)
  = (let_skel, arg_occs, scope, bind')
  where
    (rhs_skel, rhs_arg_occs, rhs') = tagSkeletonRhs bndr rhs
    arg_occs = (body_arg_occs `unionVarSet` rhs_arg_occs) `delVarSet` bndr
    bind_skel
      | is_lne    = rhs_skel -- no closure is allocated for let-no-escapes
      | otherwise = ClosureSk bndr (freeVarsOfRhs rhs) rhs_skel
    let_skel = bothSk body_skel bind_skel
    occurs_as_arg = bndr `elemVarSet` body_arg_occs
    -- Compared to the recursive case, this exploits the fact that @bndr@ is
    -- never free in @rhs@.
    scope = body_skel
    bind' = StgNonRec (BindsClosure bndr occurs_as_arg) rhs'
tagSkeletonBinding is_lne body_skel body_arg_occs (StgRec pairs)
  = (let_skel, arg_occs, scope, StgRec pairs')
  where
    (bndrs, _) = unzip pairs
    -- Local recursive STG bindings also regard the defined binders as free
    -- vars. We want to delete those for our cost model, as these are known
    -- calls anyway when we add them to the same top-level recursive group as
    -- the top-level binding currently being analysed.
    skel_occs_rhss' = map (uncurry tagSkeletonRhs) pairs
    rhss_arg_occs = map sndOf3 skel_occs_rhss'
    scope_occs = unionVarSets (body_arg_occs:rhss_arg_occs)
    arg_occs = scope_occs `delVarSetList` bndrs
    -- @skel_rhss@ aren't yet wrapped in closures. We'll do that in a moment,
    -- but we also need the un-wrapped skeletons for calculating the @scope@
    -- of the group, as the outer closures don't contribute to closure growth
    -- when we lift this specific binding.
    scope = foldr (bothSk . fstOf3) body_skel skel_occs_rhss'
    -- Now we can build the actual Skeleton for the expression just by
    -- iterating over each bind pair.
    (bind_skels, pairs') = unzip (zipWith single_bind bndrs skel_occs_rhss')
    let_skel = foldr bothSk body_skel bind_skels
    single_bind bndr (skel_rhs, _, rhs') = (bind_skel, (bndr', rhs'))
      where
        -- Here, we finally add the closure around each @skel_rhs@.
        bind_skel
          | is_lne    = skel_rhs -- no closure is allocated for let-no-escapes
          | otherwise = ClosureSk bndr fvs skel_rhs
        fvs = freeVarsOfRhs rhs' `dVarSetMinusVarSet` mkVarSet bndrs
        bndr' = BindsClosure bndr (bndr `elemVarSet` scope_occs)

tagSkeletonRhs :: Id -> CgStgRhs -> (Skeleton, IdSet, LlStgRhs)
tagSkeletonRhs _ (StgRhsCon ccs dc args)
  = (NilSk, mkArgOccs args, StgRhsCon ccs dc args)
tagSkeletonRhs bndr (StgRhsClosure fvs ccs upd bndrs body)
  = (rhs_skel, body_arg_occs, StgRhsClosure fvs ccs upd bndrs' body')
  where
    bndrs' = map BoringBinder bndrs
    (body_skel, body_arg_occs, body') = tagSkeletonExpr body
    rhs_skel = rhsSk (rhsDmdShell bndr) body_skel

-- | How many times will the lambda body of the RHS bound to the given
-- identifier be evaluated, relative to its defining context? This function
-- computes the answer in form of a 'DmdShell'.
rhsDmdShell :: Id -> DmdShell
rhsDmdShell bndr
  | is_thunk = oneifyDmd ds
  | otherwise = peelManyCalls (idArity bndr) cd
  where
    is_thunk = idArity bndr == 0
    -- Let's pray idDemandInfo is still OK after unarise...
    (ds, cd) = toCleanDmd (idDemandInfo bndr)

tagSkeletonAlt :: CgStgAlt -> (Skeleton, IdSet, LlStgAlt)
tagSkeletonAlt (con, bndrs, rhs)
  = (alt_skel, arg_occs, (con, map BoringBinder bndrs, rhs'))
  where
    (alt_skel, alt_arg_occs, rhs') = tagSkeletonExpr rhs
    arg_occs = alt_arg_occs `delVarSetList` bndrs

-- | Combines several heuristics to decide whether to lambda-lift a given
-- @let@-binding to top-level. See "StgLiftLams.Analysis#when" for details.
goodToLift
  :: DynFlags
  -> TopLevelFlag
  -> RecFlag
  -> (DIdSet -> DIdSet) -- ^ An expander function, turning 'InId's into
                        -- 'OutId's. See 'StgLiftLams.LiftM.liftedIdsExpander'.
  -> [(BinderInfo, LlStgRhs)]
  -> Skeleton
  -> Maybe DIdSet       -- ^ @Just abs_ids@ <=> This binding is beneficial to
                        -- lift and @abs_ids@ are the variables it would
                        -- abstract over
goodToLift dflags top_lvl rec_flag expander pairs scope = decide
  [ ("top-level", isTopLevel top_lvl) -- keep in sync with Note [When to lift]
  , ("memoized", any_memoized)
  , ("argument occurrences", arg_occs)
  , ("join point", is_join_point)
  , ("abstracts join points", abstracts_join_ids)
  , ("abstracts known local function", abstracts_known_local_fun)
  , ("args spill on stack", args_spill_on_stack)
  , ("increases allocation", inc_allocs)
  ] where
      decide deciders
        | not (fancy_or deciders)
        = llTrace "stgLiftLams:lifting"
                  (ppr bndrs <+> ppr abs_ids $$
                   ppr allocs $$
                   ppr scope) $
          Just abs_ids
        | otherwise
        = Nothing
      ppr_deciders = vcat . map (text . fst) . filter snd
      fancy_or deciders
        = llTrace "stgLiftLams:goodToLift" (ppr bndrs $$ ppr_deciders deciders) $
          any snd deciders

      bndrs = map (binderInfoBndr . fst) pairs
      bndrs_set = mkVarSet bndrs
      rhss = map snd pairs

      -- First objective: Calculate @abs_ids@, e.g. the former free variables
      -- the lifted binding would abstract over. We have to merge the free
      -- variables of all RHS to get the set of variables that will have to be
      -- passed through parameters.
      fvs = unionDVarSets (map freeVarsOfRhs rhss)
      -- To lift the binding to top-level, we want to delete the lifted binders
      -- themselves from the free var set. Local let bindings track recursive
      -- occurrences in their free variable set. We neither want to apply our
      -- cost model to them (see 'tagSkeletonRhs'), nor pass them as parameters
      -- when lifted, as these are known calls. We call the resulting set the
      -- identifiers we abstract over, thus @abs_ids@. These are all 'OutId's.
      -- We will save the set in 'LiftM.e_expansions' for each of the variables
      -- if we perform the lift.
      abs_ids = expander (delDVarSetList fvs bndrs)

      -- We don't lift updatable thunks or constructors
      any_memoized = any is_memoized_rhs rhss
      is_memoized_rhs StgRhsCon{} = True
      is_memoized_rhs (StgRhsClosure _ _ upd _ _) = isUpdatable upd

      -- Don't lift binders occuring as arguments. This would result in complex
      -- argument expressions which would have to be given a name, reintroducing
      -- the very allocation at each call site that we wanted to get rid off in
      -- the first place.
      arg_occs = or (mapMaybe (binderInfoOccursAsArg . fst) pairs)

      -- These don't allocate anyway.
      is_join_point = any isJoinId bndrs

      -- Abstracting over join points/let-no-escapes spoils them.
      abstracts_join_ids = any isJoinId (dVarSetElems abs_ids)

      -- Abstracting over known local functions that aren't floated themselves
      -- turns a known, fast call into an unknown, slow call:
      --
      --    let f x = ...
      --        g y = ... f x ... -- this was a known call
      --    in g 4
      --
      -- After lifting @g@, but not @f@:
      --
      --    l_g f y = ... f y ... -- this is now an unknown call
      --    let f x = ...
      --    in l_g f 4
      --
      -- We can abuse the results of arity analysis for this:
      -- idArity f > 0 ==> known
      known_fun id = idArity id > 0
      abstracts_known_local_fun
        = not (liftLamsKnown dflags) && any known_fun (dVarSetElems abs_ids)

      -- Number of arguments of a RHS in the current binding group if we decide
      -- to lift it
      n_args
        = length
        . StgCmmClosure.nonVoidIds -- void parameters don't appear in Cmm
        . (dVarSetElems abs_ids ++)
        . rhsLambdaBndrs
      max_n_args
        | isRec rec_flag = liftLamsRecArgs dflags
        | otherwise      = liftLamsNonRecArgs dflags
      -- We have 5 hardware registers on x86_64 to pass arguments in. Any excess
      -- args are passed on the stack, which means slow memory accesses
      args_spill_on_stack
        | Just n <- max_n_args = maximum (map n_args rhss) > n
        | otherwise = False

      -- We only perform the lift if allocations didn't increase.
      -- Note that @clo_growth@ will be 'infinity' if there was positive growth
      -- under a multi-shot lambda.
      -- Also, abstracting over LNEs is unacceptable. LNEs might return
      -- unlifted tuples, which idClosureFootprint can't cope with.
      inc_allocs = abstracts_join_ids || allocs > 0
      allocs = clo_growth + mkIntWithInf (negate closuresSize)
      -- We calculate and then add up the size of each binding's closure.
      -- GHC does not currently share closure environments, and we either lift
      -- the entire recursive binding group or none of it.
      closuresSize = sum $ flip map rhss $ \rhs ->
        closureSize dflags
        . dVarSetElems
        . expander
        . flip dVarSetMinusVarSet bndrs_set
        $ freeVarsOfRhs rhs
      clo_growth = closureGrowth expander (idClosureFootprint dflags) bndrs_set abs_ids scope

rhsLambdaBndrs :: LlStgRhs -> [Id]
rhsLambdaBndrs StgRhsCon{} = []
rhsLambdaBndrs (StgRhsClosure _ _ _ bndrs _) = map binderInfoBndr bndrs

-- | The size in words of a function closure closing over the given 'Id's,
-- including the header.
closureSize :: DynFlags -> [Id] -> WordOff
closureSize dflags ids = words + sTD_HDR_SIZE dflags
  -- We go through sTD_HDR_SIZE rather than fixedHdrSizeW so that we don't
  -- optimise differently when profiling is enabled.
  where
    (words, _, _)
      -- Functions have a StdHeader (as opposed to ThunkHeader).
      = StgCmmLayout.mkVirtHeapOffsets dflags StgCmmLayout.StdHeader
      . StgCmmClosure.addIdReps
      . StgCmmClosure.nonVoidIds
      $ ids

-- | The number of words a single 'Id' adds to a closure's size.
-- Note that this can't handle unboxed tuples (which may still be present in
-- let-no-escapes, even after Unarise), in which case
-- @'StgCmmClosure.idPrimRep'@ will crash.
idClosureFootprint:: DynFlags -> Id -> WordOff
idClosureFootprint dflags
  = StgCmmArgRep.argRepSizeW dflags
  . StgCmmArgRep.idArgRep

-- | @closureGrowth expander sizer f fvs@ computes the closure growth in words
-- as a result of lifting @f@ to top-level. If there was any growing closure
-- under a multi-shot lambda, the result will be 'infinity'.
-- Also see "StgLiftLams.Analysis#clogro".
closureGrowth
  :: (DIdSet -> DIdSet)
  -- ^ Expands outer free ids that were lifted to their free vars
  -> (Id -> Int)
  -- ^ Computes the closure footprint of an identifier
  -> IdSet
  -- ^ Binding group for which lifting is to be decided
  -> DIdSet
  -- ^ Free vars of the whole binding group prior to lifting it. These must be
  --   available at call sites if we decide to lift the binding group.
  -> Skeleton
  -- ^ Abstraction of the scope of the function
  -> IntWithInf
  -- ^ Closure growth. 'infinity' indicates there was growth under a
  --   (multi-shot) lambda.
closureGrowth expander sizer group abs_ids = go
  where
    go NilSk = 0
    go (BothSk a b) = go a + go b
    go (AltSk a b) = max (go a) (go b)
    go (ClosureSk _ clo_fvs rhs)
      -- If no binder of the @group@ occurs free in the closure, the lifting
      -- won't have any effect on it and we can omit the recursive call.
      | n_occs == 0 = 0
      -- Otherwise, we account the cost of allocating the closure and add it to
      -- the closure growth of its RHS.
      | otherwise   = mkIntWithInf cost + go rhs
      where
        n_occs = sizeDVarSet (clo_fvs' `dVarSetIntersectVarSet` group)
        -- What we close over considering prior lifting decisions
        clo_fvs' = expander clo_fvs
        -- Variables that would additionally occur free in the closure body if
        -- we lift @f@
        newbies = abs_ids `minusDVarSet` clo_fvs'
        -- Lifting @f@ removes @f@ from the closure but adds all @newbies@
        cost = foldDVarSet (\id size -> sizer id + size) 0 newbies - n_occs
    go (RhsSk body_dmd body)
      -- The conservative assumption would be that
      --   1. Every RHS with positive growth would be called multiple times,
      --      modulo thunks.
      --   2. Every RHS with negative growth wouldn't be called at all.
      --
      -- In the first case, we'd have to return 'infinity', while in the
      -- second case, we'd have to return 0. But we can do far better
      -- considering information from the demand analyser, which provides us
      -- with conservative estimates on minimum and maximum evaluation
      -- cardinality. The @body_dmd@ part of 'RhsSk' is the result of
      -- 'rhsDmdShell' and accurately captures the cardinality of the RHSs body
      -- relative to its defining context.
      | isAbsDmd body_dmd   = 0
      | cg <= 0             = if isStrictDmd body_dmd then cg else 0
      | isUsedOnce body_dmd = cg
      | otherwise           = infinity
      where
        cg = go body