module SAT ( doStaticArgs ) where
import GhcPrelude
import Var
import CoreSyn
import CoreUtils
import Type
import Coercion
import Id
import Name
import VarEnv
import UniqSupply
import Util
import UniqFM
import VarSet
import Unique
import UniqSet
import Outputable
import Data.List
import FastString
#include "HsVersions.h"
doStaticArgs :: UniqSupply -> CoreProgram -> CoreProgram
doStaticArgs us binds = snd $ mapAccumL sat_bind_threaded_us us binds
where
sat_bind_threaded_us us bind =
let (us1, us2) = splitUniqSupply us
in (us1, fst $ runSAT us2 (satBind bind emptyUniqSet))
satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
satBind (NonRec binder expr) interesting_ids = do
(expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
return (NonRec binder expr', finalizeApp expr_app sat_info_expr)
satBind (Rec [(binder, rhs)]) interesting_ids = do
let interesting_ids' = interesting_ids `addOneToUniqSet` binder
(rhs_binders, rhs_body) = collectBinders rhs
(rhs_body', sat_info_rhs_body) <- satTopLevelExpr rhs_body interesting_ids'
let sat_info_rhs_from_args = unitVarEnv binder (bindersToSATInfo rhs_binders)
sat_info_rhs' = mergeIdSATInfo sat_info_rhs_from_args sat_info_rhs_body
shadowing = binder `elementOfUniqSet` interesting_ids
sat_info_rhs'' = if shadowing
then sat_info_rhs' `delFromUFM` binder
else sat_info_rhs'
bind' <- saTransformMaybe binder (lookupUFM sat_info_rhs' binder)
rhs_binders rhs_body'
return (bind', sat_info_rhs'')
satBind (Rec pairs) interesting_ids = do
let (binders, rhss) = unzip pairs
rhss_SATed <- mapM (\e -> satTopLevelExpr e interesting_ids) rhss
let (rhss', sat_info_rhss') = unzip rhss_SATed
return (Rec (zipEqual "satBind" binders rhss'), mergeIdSATInfos sat_info_rhss')
data App = VarApp Id | TypeApp Type | CoApp Coercion
data Staticness a = Static a | NotStatic
type IdAppInfo = (Id, SATInfo)
type SATInfo = [Staticness App]
type IdSATInfo = IdEnv SATInfo
emptyIdSATInfo :: IdSATInfo
emptyIdSATInfo = emptyUFM
pprSATInfo :: SATInfo -> SDoc
pprSATInfo staticness = hcat $ map pprStaticness staticness
pprStaticness :: Staticness App -> SDoc
pprStaticness (Static (VarApp _)) = text "SV"
pprStaticness (Static (TypeApp _)) = text "ST"
pprStaticness (Static (CoApp _)) = text "SC"
pprStaticness NotStatic = text "NS"
mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
mergeSATInfo l r = zipWith mergeSA l r
where
mergeSA NotStatic _ = NotStatic
mergeSA _ NotStatic = NotStatic
mergeSA (Static (VarApp v)) (Static (VarApp v'))
| v == v' = Static (VarApp v)
| otherwise = NotStatic
mergeSA (Static (TypeApp t)) (Static (TypeApp t'))
| t `eqType` t' = Static (TypeApp t)
| otherwise = NotStatic
mergeSA (Static (CoApp c)) (Static (CoApp c'))
| c `eqCoercion` c' = Static (CoApp c)
| otherwise = NotStatic
mergeSA _ _ = pprPanic "mergeSATInfo" $
text "Left:"
<> pprSATInfo l <> text ", "
<> text "Right:"
<> pprSATInfo r
mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo = plusUFM_C mergeSATInfo
mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
mergeIdSATInfos = foldl' mergeIdSATInfo emptyIdSATInfo
bindersToSATInfo :: [Id] -> SATInfo
bindersToSATInfo vs = map (Static . binderToApp) vs
where binderToApp v | isId v = VarApp v
| isTyVar v = TypeApp $ mkTyVarTy v
| otherwise = CoApp $ mkCoVarCo v
finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Nothing id_sat_info = id_sat_info
finalizeApp (Just (v, sat_info')) id_sat_info =
let sat_info'' = case lookupUFM id_sat_info v of
Nothing -> sat_info'
Just sat_info -> mergeSATInfo sat_info sat_info'
in extendVarEnv id_sat_info v sat_info''
satTopLevelExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo)
satTopLevelExpr expr interesting_ids = do
(expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
return (expr', finalizeApp expr_app sat_info_expr)
satExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
satExpr var@(Var v) interesting_ids = do
let app_info = if v `elementOfUniqSet` interesting_ids
then Just (v, [])
else Nothing
return (var, emptyIdSATInfo, app_info)
satExpr lit@(Lit _) _ = do
return (lit, emptyIdSATInfo, Nothing)
satExpr (Lam binders body) interesting_ids = do
(body', sat_info, this_app) <- satExpr body interesting_ids
return (Lam binders body', finalizeApp this_app sat_info, Nothing)
satExpr (App fn arg) interesting_ids = do
(fn', sat_info_fn, fn_app) <- satExpr fn interesting_ids
let satRemainder = boring fn' sat_info_fn
case fn_app of
Nothing -> satRemainder Nothing
Just (fn_id, fn_app_info) ->
let satRemainderWithStaticness arg_staticness = satRemainder $ Just (fn_id, fn_app_info ++ [arg_staticness])
in case arg of
Type t -> satRemainderWithStaticness $ Static (TypeApp t)
Coercion c -> satRemainderWithStaticness $ Static (CoApp c)
Var v -> satRemainderWithStaticness $ Static (VarApp v)
_ -> satRemainderWithStaticness $ NotStatic
where
boring :: CoreExpr -> IdSATInfo -> Maybe IdAppInfo -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
boring fn' sat_info_fn app_info =
do (arg', sat_info_arg, arg_app) <- satExpr arg interesting_ids
let sat_info_arg' = finalizeApp arg_app sat_info_arg
sat_info = mergeIdSATInfo sat_info_fn sat_info_arg'
return (App fn' arg', sat_info, app_info)
satExpr (Case expr bndr ty alts) interesting_ids = do
(expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
let sat_info_expr' = finalizeApp expr_app sat_info_expr
zipped_alts' <- mapM satAlt alts
let (alts', sat_infos_alts) = unzip zipped_alts'
return (Case expr' bndr ty alts', mergeIdSATInfo sat_info_expr' (mergeIdSATInfos sat_infos_alts), Nothing)
where
satAlt (con, bndrs, expr) = do
(expr', sat_info_expr) <- satTopLevelExpr expr interesting_ids
return ((con, bndrs, expr'), sat_info_expr)
satExpr (Let bind body) interesting_ids = do
(body', sat_info_body, body_app) <- satExpr body interesting_ids
(bind', sat_info_bind) <- satBind bind interesting_ids
return (Let bind' body', mergeIdSATInfo sat_info_body sat_info_bind, body_app)
satExpr (Tick tickish expr) interesting_ids = do
(expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
return (Tick tickish expr', sat_info_expr, expr_app)
satExpr ty@(Type _) _ = do
return (ty, emptyIdSATInfo, Nothing)
satExpr co@(Coercion _) _ = do
return (co, emptyIdSATInfo, Nothing)
satExpr (Cast expr coercion) interesting_ids = do
(expr', sat_info_expr, expr_app) <- satExpr expr interesting_ids
return (Cast expr' coercion, sat_info_expr, expr_app)
type SatM result = UniqSM result
runSAT :: UniqSupply -> SatM a -> a
runSAT = initUs_
newUnique :: SatM Unique
newUnique = getUniqueM
saTransformMaybe :: Id -> Maybe SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
saTransformMaybe binder maybe_arg_staticness rhs_binders rhs_body
| Just arg_staticness <- maybe_arg_staticness
, should_transform arg_staticness
= saTransform binder arg_staticness rhs_binders rhs_body
| otherwise
= return (Rec [(binder, mkLams rhs_binders rhs_body)])
where
should_transform staticness = n_static_args > 1
where
n_static_args = count isStaticValue staticness
saTransform :: Id -> SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
saTransform binder arg_staticness rhs_binders rhs_body
= do { shadow_lam_bndrs <- mapM clone binders_w_staticness
; uniq <- newUnique
; return (NonRec binder (mk_new_rhs uniq shadow_lam_bndrs)) }
where
binders_w_staticness = rhs_binders `zip` (arg_staticness ++ repeat NotStatic)
non_static_args :: [Var]
non_static_args = [v | (v, NotStatic) <- binders_w_staticness]
clone (bndr, NotStatic) = return bndr
clone (bndr, _ ) = do { uniq <- newUnique
; return (setVarUnique bndr uniq) }
mk_new_rhs uniq shadow_lam_bndrs
= mkLams rhs_binders $
Let (Rec [(rec_body_bndr, rec_body)])
local_body
where
local_body = mkVarApps (Var rec_body_bndr) non_static_args
rec_body = mkLams non_static_args $
Let (NonRec shadow_bndr shadow_rhs) rhs_body
shadow_rhs = mkLams shadow_lam_bndrs local_body
rec_body_bndr = mkSysLocal (fsLit "sat_worker") uniq (exprType rec_body)
shadow_bndr = mkSysLocal (occNameFS (getOccName binder))
(idUnique binder)
(exprType shadow_rhs)
isStaticValue :: Staticness App -> Bool
isStaticValue (Static (VarApp _)) = True
isStaticValue _ = False