module GHC.Stg.BcPrep ( bcPrep ) where
import GHC.Prelude
import GHC.Types.Id.Make
import GHC.Types.Id
import GHC.Core.Type
import GHC.Builtin.Types ( unboxedUnitTy )
import GHC.Builtin.Types.Prim
import GHC.Types.Unique
import GHC.Data.FastString
import GHC.Utils.Panic.Plain
import GHC.Types.Tickish
import GHC.Types.Unique.Supply
import qualified GHC.Types.CostCentre as CC
import GHC.Stg.Syntax
import GHC.Utils.Monad.State.Strict
data BcPrepM_State
= BcPrepM_State
{ prepUniqSupply :: !UniqSupply
}
type BcPrepM a = State BcPrepM_State a
bcPrepRHS :: StgRhs -> BcPrepM StgRhs
bcPrepRHS (StgRhsClosure fvs cc upd args (StgTick bp@Breakpoint{} expr)) = do
expr' <- bcPrepExpr expr
pure (StgRhsClosure fvs cc upd args (StgTick bp expr'))
bcPrepRHS (StgRhsClosure fvs cc upd args expr) =
StgRhsClosure fvs cc upd args <$> bcPrepExpr expr
bcPrepRHS con@StgRhsCon{} = pure con
bcPrepExpr :: StgExpr -> BcPrepM StgExpr
bcPrepExpr (StgTick bp@(Breakpoint tick_ty _ _) rhs)
| isLiftedTypeKind (typeKind tick_ty) = do
id <- newId tick_ty
rhs' <- bcPrepExpr rhs
let expr' = StgTick bp rhs'
bnd = StgNonRec id (StgRhsClosure noExtFieldSilent
CC.dontCareCCS
ReEntrant
[]
expr'
)
letExp = StgLet noExtFieldSilent bnd (StgApp id [])
pure letExp
| otherwise = do
id <- newId (mkVisFunTyMany realWorldStatePrimTy tick_ty)
rhs' <- bcPrepExpr rhs
let expr' = StgTick bp rhs'
bnd = StgNonRec id (StgRhsClosure noExtFieldSilent
CC.dontCareCCS
ReEntrant
[voidArgId]
expr'
)
pure $ StgLet noExtFieldSilent bnd (StgApp id [StgVarArg realWorldPrimId])
bcPrepExpr (StgTick tick rhs) =
StgTick tick <$> bcPrepExpr rhs
bcPrepExpr (StgLet xlet bnds expr) =
StgLet xlet <$> bcPrepBind bnds
<*> bcPrepExpr expr
bcPrepExpr (StgLetNoEscape xlne bnds expr) =
StgLet xlne <$> bcPrepBind bnds
<*> bcPrepExpr expr
bcPrepExpr (StgCase expr bndr alt_type alts) =
StgCase <$> bcPrepExpr expr
<*> pure bndr
<*> pure alt_type
<*> mapM bcPrepAlt alts
bcPrepExpr lit@StgLit{} = pure lit
bcPrepExpr (StgApp x [])
| isNNLJoinPoint x = pure $
StgApp (protectNNLJoinPointId x) [StgVarArg voidPrimId]
bcPrepExpr app@StgApp{} = pure app
bcPrepExpr app@StgConApp{} = pure app
bcPrepExpr app@StgOpApp{} = pure app
bcPrepAlt :: StgAlt -> BcPrepM StgAlt
bcPrepAlt (GenStgAlt con bndrs rhs) = GenStgAlt con bndrs <$> bcPrepExpr rhs
bcPrepBind :: StgBinding -> BcPrepM StgBinding
bcPrepBind (StgNonRec bndr rhs) =
let (bndr', rhs') = bcPrepSingleBind (bndr, rhs)
in StgNonRec bndr' <$> bcPrepRHS rhs'
bcPrepBind (StgRec bnds) =
StgRec <$> mapM ((\(b,r) -> (,) b <$> bcPrepRHS r) . bcPrepSingleBind)
bnds
bcPrepSingleBind :: (Id, StgRhs) -> (Id, StgRhs)
bcPrepSingleBind (x, StgRhsClosure ext cc upd_flag args body)
| isNNLJoinPoint x
= ( protectNNLJoinPointId x
, StgRhsClosure ext cc upd_flag (args ++ [voidArgId]) body)
bcPrepSingleBind bnd = bnd
bcPrepTopLvl :: StgTopBinding -> BcPrepM StgTopBinding
bcPrepTopLvl lit@StgTopStringLit{} = pure lit
bcPrepTopLvl (StgTopLifted bnd) = StgTopLifted <$> bcPrepBind bnd
bcPrep :: UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
bcPrep us bnds = evalState (mapM bcPrepTopLvl bnds) (BcPrepM_State us)
isNNLJoinPoint :: Id -> Bool
isNNLJoinPoint x = isJoinId x && mightBeUnliftedType (idType x)
protectNNLJoinPointId :: Id -> Id
protectNNLJoinPointId x
= assert (isNNLJoinPoint x )
updateIdTypeButNotMult (unboxedUnitTy `mkVisFunTyMany`) x
newUnique :: BcPrepM Unique
newUnique = state $
\st -> case takeUniqFromSupply (prepUniqSupply st) of
(uniq, us) -> (uniq, st { prepUniqSupply = us })
newId :: Type -> BcPrepM Id
newId ty = do
uniq <- newUnique
return $ mkSysLocal prepFS uniq Many ty
prepFS :: FastString
prepFS = fsLit "bcprep"