{-# LANGUAGE TypeFamilies, DataKinds, GADTs, FlexibleInstances #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}

-- This module declares some basic types used by GHC.Stg.Lift
-- We can import this module into GHC.Stg.Syntax, where the
-- type instance declarations for BinderP etc live

module GHC.Stg.Lift.Types(
   Skeleton(..),
   bothSk, altSk, rhsSk,

   BinderInfo(..),
   binderInfoBndr, binderInfoOccursAsArg
   ) where

import GHC.Prelude

import GHC.Types.Id
import GHC.Types.Demand
import GHC.Types.Var.Set

import GHC.Utils.Outputable

-- | 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 !Card {- ^ how often the RHS was entered -} !Skeleton
  | AltSk !Skeleton !Skeleton
  | BothSk !Skeleton !Skeleton
  | NilSk

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

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

rhsSk :: Card -> Skeleton -> Skeleton
rhsSk :: Card -> Skeleton -> Skeleton
rhsSk Card
_        Skeleton
NilSk = Skeleton
NilSk
rhsSk Card
body_dmd Skeleton
skel  = Card -> Skeleton -> Skeleton
RhsSk Card
body_dmd Skeleton
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 "GHC.Stg.Lift.Analysis#arg_occs").
  | BoringBinder !Id       -- ^ Every other kind of binder

-- | Gets the bound 'Id' out a 'BinderInfo'.
binderInfoBndr :: BinderInfo -> Id
binderInfoBndr :: BinderInfo -> Id
binderInfoBndr (BoringBinder Id
bndr)   = Id
bndr
binderInfoBndr (BindsClosure Id
bndr Bool
_) = Id
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 :: BinderInfo -> Maybe Bool
binderInfoOccursAsArg BoringBinder{}     = Maybe Bool
forall a. Maybe a
Nothing
binderInfoOccursAsArg (BindsClosure Id
_ Bool
b) = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
b

instance Outputable Skeleton where
  ppr :: Skeleton -> SDoc
ppr Skeleton
NilSk = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
""
  ppr (AltSk Skeleton
l Skeleton
r) = [SDoc] -> SDoc
forall doc. IsDoc doc => [doc] -> doc
vcat
    [ String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"{ " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
l
    , String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ALT"
    , String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"  " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
r
    , String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"}"
    ]
  ppr (BothSk Skeleton
l Skeleton
r) = Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
l SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
r
  ppr (ClosureSk Id
f DIdSet
fvs Skeleton
body) = Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr Id
f SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> DIdSet -> SDoc
forall a. Outputable a => a -> SDoc
ppr DIdSet
fvs SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ Int -> SDoc -> SDoc
nest Int
2 (Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
body)
  ppr (RhsSk Card
card Skeleton
body) = [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat
    [ SDoc
lambda
    , Card -> SDoc
forall a. Outputable a => a -> SDoc
ppr Card
card
    , SDoc
forall doc. IsLine doc => doc
dot
    , Skeleton -> SDoc
forall a. Outputable a => a -> SDoc
ppr Skeleton
body
    ]

instance Outputable BinderInfo where
  ppr :: BinderInfo -> SDoc
ppr = Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr (Id -> SDoc) -> (BinderInfo -> Id) -> BinderInfo -> SDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderInfo -> Id
binderInfoBndr

instance OutputableBndr BinderInfo where
  pprBndr :: BindingSite -> BinderInfo -> SDoc
pprBndr BindingSite
b = BindingSite -> Id -> SDoc
forall a. OutputableBndr a => BindingSite -> a -> SDoc
pprBndr BindingSite
b (Id -> SDoc) -> (BinderInfo -> Id) -> BinderInfo -> SDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderInfo -> Id
binderInfoBndr
  pprPrefixOcc :: BinderInfo -> SDoc
pprPrefixOcc = Id -> SDoc
forall a. OutputableBndr a => a -> SDoc
pprPrefixOcc (Id -> SDoc) -> (BinderInfo -> Id) -> BinderInfo -> SDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderInfo -> Id
binderInfoBndr
  pprInfixOcc :: BinderInfo -> SDoc
pprInfixOcc = Id -> SDoc
forall a. OutputableBndr a => a -> SDoc
pprInfixOcc (Id -> SDoc) -> (BinderInfo -> Id) -> BinderInfo -> SDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderInfo -> Id
binderInfoBndr
  bndrIsJoin_maybe :: BinderInfo -> JoinPointHood
bndrIsJoin_maybe = Id -> JoinPointHood
forall a. OutputableBndr a => a -> JoinPointHood
bndrIsJoin_maybe (Id -> JoinPointHood)
-> (BinderInfo -> Id) -> BinderInfo -> JoinPointHood
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderInfo -> Id
binderInfoBndr