import CoreSyn
import CoreLint ( showPass, endPass )
-import CoreUtils ( exprType, exprIsConApp_maybe, eqExpr )
+import CoreUtils ( exprType, tcEqExpr, mkPiTypes )
import CoreFVs ( exprsFreeVars )
-import DataCon ( isExistentialDataCon )
-import PprCore ( pprCoreRules )
-import Id ( Id, idName, idSpecialisation, mkUserLocal, mkSysLocal )
+import CoreSubst ( Subst, mkSubst, substExpr )
+import CoreTidy ( tidyRules )
+import PprCore ( pprRules )
+import WwLib ( mkWorkerArgs )
+import DataCon ( dataConRepArity, isVanillaDataCon )
+import Type ( tyConAppArgs, tyVarsOfTypes )
+import Unify ( coreRefineTys )
+import Id ( Id, idName, idType, isDataConWorkId_maybe,
+ mkUserLocal, mkSysLocal )
import Var ( Var )
import VarEnv
import VarSet
import Name ( nameOccName, nameSrcLoc )
-import Rules ( addIdSpecialisations )
+import Rules ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
import OccName ( mkSpecOcc )
import ErrUtils ( dumpIfSet_dyn )
-import CmdLineOpts ( DynFlags, DynFlag(..) )
-import Outputable
-
+import DynFlags ( DynFlags, DynFlag(..) )
+import BasicTypes ( Activation(..) )
import Maybes ( orElse )
-import Util ( mapAccumL )
+import Util ( mapAccumL, lengthAtLeast, notNull )
import List ( nubBy, partition )
import UniqSupply
import Outputable
+import FastString
\end{code}
-----------------------------------------------------
endPass dflags "SpecConstr" Opt_D_dump_spec binds'
dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
- (vcat (map dump_specs (concat (map bindersOf binds'))))
+ (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
return binds'
where
go env (bind:binds) = scBind env bind `thenUs` \ (env', _, bind') ->
go env' binds `thenUs` \ binds' ->
returnUs (bind' : binds')
-
-dump_specs var = pprCoreRules var (idSpecialisation var)
\end{code}
%************************************************************************
%* *
-\subsection{Environments and such}
+\subsection{Environment: goes downwards}
%* *
%************************************************************************
\begin{code}
-type ScEnv = VarEnv HowBound
+data ScEnv = SCE { scope :: VarEnv HowBound,
+ -- Binds all non-top-level variables in scope
-emptyScEnv = emptyVarEnv
+ cons :: ConstrEnv
+ }
+
+type ConstrEnv = IdEnv ConValue
+data ConValue = CV AltCon [CoreArg]
+ -- Variables known to be bound to a constructor
+ -- in a particular case alternative
+
+refineConstrEnv :: Subst -> ConstrEnv -> ConstrEnv
+-- The substitution is a type substitution only
+refineConstrEnv subst env = mapVarEnv refine_con_value env
+ where
+ refine_con_value (CV con args) = CV con (map (substExpr subst) args)
+
+emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
data HowBound = RecFun -- These are the recursive functions for which
-- we seek interesting call patterns
+
| RecArg -- These are those functions' arguments; we are
-- interested to see if those arguments are scrutinised
+
| Other -- We track all others so we know what's in scope
+ -- This is used in spec_one to check what needs to be
+ -- passed as a parameter and what is in scope at the
+ -- function definition site
+
+instance Outputable HowBound where
+ ppr RecFun = text "RecFun"
+ ppr RecArg = text "RecArg"
+ ppr Other = text "Other"
+
+lookupScopeEnv env v = lookupVarEnv (scope env) v
+
+extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
+extendBndr env bndr = env { scope = extendVarEnv (scope env) bndr Other }
+
+ -- When we encounter
+ -- case scrut of b
+ -- C x y -> ...
+ -- we want to bind b, and perhaps scrut too, to (C x y)
+extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
+extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
+ = extendBndrs env (case_bndr : alt_bndrs)
+
+extendCaseBndrs env case_bndr scrut con@(LitAlt lit) alt_bndrs
+ = ASSERT( null alt_bndrs ) extendAlt env case_bndr scrut (CV con []) []
+
+extendCaseBndrs env case_bndr scrut con@(DataAlt data_con) alt_bndrs
+ | isVanillaDataCon data_con
+ = extendAlt env case_bndr scrut (CV con vanilla_args) alt_bndrs
+
+ | otherwise -- GADT
+ = extendAlt env1 case_bndr scrut (CV con gadt_args) alt_bndrs
+ where
+ vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
+ map varToCoreExpr alt_bndrs
+
+ gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
+
+ (alt_tvs, _) = span isTyVar alt_bndrs
+ Just (tv_subst, is_local) = coreRefineTys data_con alt_tvs (idType case_bndr)
+ subst = mkSubst in_scope tv_subst emptyVarEnv -- No Id substitition
+ in_scope = mkInScopeSet (tyVarsOfTypes (varEnvElts tv_subst))
+
+ env1 | is_local = env
+ | otherwise = env { cons = refineConstrEnv subst (cons env) }
+
+
+
+extendAlt :: ScEnv -> Id -> CoreExpr -> ConValue -> [Var] -> ScEnv
+extendAlt env case_bndr scrut val alt_bndrs
+ = let
+ env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
+ cons = extendVarEnv (cons env) case_bndr val }
+ in
+ case scrut of
+ Var v -> -- Bind the scrutinee in the ConstrEnv if it's a variable
+ -- Also forget if the scrutinee is a RecArg, because we're
+ -- now in the branch of a case, and we don't want to
+ -- record a non-scrutinee use of v if we have
+ -- case v of { (a,b) -> ...(f v)... }
+ SCE { scope = extendVarEnv (scope env1) v Other,
+ cons = extendVarEnv (cons env1) v val }
+ other -> env1
+
+ -- When we encounter a recursive function binding
+ -- f = \x y -> ...
+ -- we want to extend the scope env with bindings
+ -- that record that f is a RecFn and x,y are RecArgs
+extendRecBndr env fn bndrs
+ = env { scope = scope env `extendVarEnvList`
+ ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
+\end{code}
-extendBndrs env bndrs = extendVarEnvList env [(b,Other) | b <- bndrs]
-extendBndr env bndr = extendVarEnv env bndr Other
+%************************************************************************
+%* *
+\subsection{Usage information: flows upwards}
+%* *
+%************************************************************************
+
+\begin{code}
data ScUsage
= SCU {
- calls :: !(IdEnv ([[CoreArg]])), -- Calls
- -- The functions are a subset of the
- -- RecFuns in the ScEnv
+ calls :: !(IdEnv ([Call])), -- Calls
+ -- The functions are a subset of the
+ -- RecFuns in the ScEnv
occs :: !(IdEnv ArgOcc) -- Information on argument occurrences
} -- The variables are a subset of the
-- RecArg in the ScEnv
+type Call = (ConstrEnv, [CoreArg])
+ -- The arguments of the call, together with the
+ -- env giving the constructor bindings at the call site
+
nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
%* *
%************************************************************************
+The main recursive function gathers up usage information, and
+creates specialised versions of functions.
+
\begin{code}
scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
-- The unique supply is needed when we invent
scExpr env (Lam b e) = scExpr (extendBndr env b) e `thenUs` \ (usg,e') ->
returnUs (usg, Lam b e')
-scExpr env (Case scrut b alts)
+scExpr env (Case scrut b ty alts)
= sc_scrut scrut `thenUs` \ (scrut_usg, scrut') ->
mapAndUnzipUs sc_alt alts `thenUs` \ (alts_usgs, alts') ->
returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
- Case scrut' b alts')
+ Case scrut' b ty alts')
where
sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
sc_scrut e = scExpr env e
- sc_alt (con,bs,rhs) = scExpr env rhs `thenUs` \ (usg,rhs') ->
+ sc_alt (con,bs,rhs) = scExpr env1 rhs `thenUs` \ (usg,rhs') ->
returnUs (usg, (con,bs,rhs'))
where
- env1 = extendBndrs env (b:bs)
+ env1 = extendCaseBndrs env b scrut con bs
scExpr env (Let bind body)
= scBind env bind `thenUs` \ (env', bind_usg, bind') ->
let
arg_usg = combineUsages usgs
fn_usg | Var f <- fn,
- Just RecFun <- lookupVarEnv env f
- = SCU { calls = unitVarEnv f [args], occs = emptyVarEnv }
+ Just RecFun <- lookupScopeEnv env f
+ = SCU { calls = unitVarEnv f [(cons env, args)],
+ occs = emptyVarEnv }
| otherwise
= nullUsage
in
----------------------
scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
scBind env (Rec [(fn,rhs)])
- | not (null val_bndrs)
- = scExpr env' body `thenUs` \ (usg@(SCU { calls = calls, occs = occs }), body') ->
+ | notNull val_bndrs
+ = scExpr env_fn_body body `thenUs` \ (usg, body') ->
+ let
+ SCU { calls = calls, occs = occs } = usg
+ in
specialise env fn bndrs body usg `thenUs` \ (rules, spec_prs) ->
- returnUs (extendBndrs env bndrs,
+ returnUs (extendBndr env fn, -- For the body of the letrec, just
+ -- extend the env with Other to record
+ -- that it's in scope; no funny RecFun business
SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
where
(bndrs,body) = collectBinders rhs
val_bndrs = filter isId bndrs
- env' = env `extendVarEnvList` ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs])
+ env_fn_body = extendRecBndr env fn bndrs
scBind env (Rec prs)
= mapAndUnzipUs do_one prs `thenUs` \ (usgs, prs') ->
----------------------
varUsage env v use
- | Just RecArg <- lookupVarEnv env v = SCU { calls = emptyVarEnv, occs = unitVarEnv v use }
- | otherwise = nullUsage
+ | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
+ occs = unitVarEnv v use }
+ | otherwise = nullUsage
\end{code}
good_calls :: [[CoreArg]]
good_calls = [ pats
- | call_args <- all_calls,
- length call_args >= n_bndrs, -- App is saturated
+ | (con_env, call_args) <- all_calls,
+ call_args `lengthAtLeast` n_bndrs, -- App is saturated
let call = (bndrs `zip` call_args),
- any (good_arg occs) call,
- let (_, pats) = argsToPats us call_args
+ any (good_arg con_env occs) call, -- At least one arg is a constr app
+ let (_, pats) = argsToPats con_env us call_args
]
in
- pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
(nubBy same_call good_calls `zip` [1..])
where
n_bndrs = length bndrs
- same_call as1 as2 = and (zipWith eqExpr as1 as2)
+ same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
---------------------
-good_arg :: IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
-good_arg arg_occs (bndr, arg)
- = case exprIsConApp_maybe arg of -- exprIsConApp_maybe looks
- Just (dc,_) -> not (isExistentialDataCon dc) -- through unfoldings
- && bndr_usg_ok arg_occs bndr arg
+good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
+good_arg con_env arg_occs (bndr, arg)
+ = case is_con_app_maybe con_env arg of
+ Just _ -> bndr_usg_ok arg_occs bndr arg
other -> False
bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
bndr_usg_ok arg_occs bndr arg
- = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
- case lookupVarEnv arg_occs bndr of
+ = case lookupVarEnv arg_occs bndr of
Just CaseScrut -> True -- Used only by case scrutiny
Just Both -> case arg of -- Used by case and elsewhere
App _ _ -> True -- so the arg should be an explicit con app
---------------------
-argsToPats :: UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
-argsToPats us args = mapAccumL argToPat us args
-
-argToPat :: UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
--- C a (D (f x) (g y)) ==> C p1 (D p2 p3)
-argToPat us (Type ty)
- = (us, Type ty)
-
-argToPat us arg
- | Just (dc,args) <- exprIsConApp_maybe arg
- = let
- (us',args') = argsToPats us args
- in
- (us', mkConApp dc args')
-
-argToPat us (Var v) -- Don't uniqify existing vars,
- = (us, Var v) -- so that we can spot when we pass them twice
-
-argToPat us arg
- = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
- where
- (us1,us2) = splitUniqSupply us
-
----------------------
spec_one :: ScEnv
-> Id -- Function
-> CoreExpr -- Rhs of the original function
-> ([CoreArg], Int)
-> UniqSM (CoreRule, (Id,CoreExpr)) -- Rule and binding
+-- spec_one creates a specialised copy of the function, together
+-- with a rule for using it. I'm very proud of how short this
+-- function is, considering what it does :-).
+
{-
Example
In-scope: a, x::a
- f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
- [c is presumably bound by the (...) part]
+ f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
+ [c::*, v::(b,c) are presumably bound by the (...) part]
==>
- f_spec = /\ b c \ v::(a,(b,c)) ->
- (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
+ f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
+ (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
- RULE: forall b c,
- y::[(a,(b,c))],
- v::(a,(b,c)),
- h::(a,(b,c))->[(a,(b,c))] .
+ RULE: forall b::* c::*, -- Note, *not* forall a, x
+ v::(b,c),
+ hw::[(a,(b,c))] .
- f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
+ f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
-}
-spec_one env fn rhs (pats, n)
- = getUniqueUs `thenUs` \ spec_uniq ->
+spec_one env fn rhs (pats, rule_number)
+ = getUniqueUs `thenUs` \ spec_uniq ->
let
fn_name = idName fn
fn_loc = nameSrcLoc fn_name
spec_occ = mkSpecOcc (nameOccName fn_name)
pat_fvs = varSetElems (exprsFreeVars pats)
vars_to_bind = filter not_avail pat_fvs
- not_avail v = not (v `elemVarEnv` env)
- -- Put the type variables first just for tidiness
+ not_avail v = not (v `elemVarEnv` scope env)
+ -- Put the type variables first; the type of a term
+ -- variable may mention a type variable
(tvs, ids) = partition isTyVar vars_to_bind
bndrs = tvs ++ ids
+ spec_body = mkApps rhs pats
+ body_ty = exprType spec_body
- rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
- spec_rhs = mkLams bndrs (mkApps rhs pats)
- spec_id = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
- rule = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
+ (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
+ -- Usual w/w hack to avoid generating
+ -- a spec_rhs of unlifted type and no args
+
+ rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
+ spec_rhs = mkLams spec_lam_args spec_body
+ spec_id = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
+ rule_rhs = mkVarApps (Var spec_id) spec_call_args
+ rule = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
in
returnUs (rule, (spec_id, spec_rhs))
+
+-- In which phase should the specialise-constructor rules be active?
+-- Originally I made them always-active, but Manuel found that
+-- this defeated some clever user-written rules. So Plan B
+-- is to make them active only in Phase 0; after all, currently,
+-- the specConstr transformation is only run after the simplifier
+-- has reached Phase 0. In general one would want it to be
+-- flag-controllable, but for now I'm leaving it baked in
+-- [SLPJ Oct 01]
+specConstrActivation :: Activation
+specConstrActivation = ActiveAfter 0 -- Baked in; see comments above
+\end{code}
+
+%************************************************************************
+%* *
+\subsection{Argument analysis}
+%* *
+%************************************************************************
+
+This code deals with analysing call-site arguments to see whether
+they are constructor applications.
+
+\begin{code}
+ -- argToPat takes an actual argument, and returns an abstracted
+ -- version, consisting of just the "constructor skeleton" of the
+ -- argument, with non-constructor sub-expression replaced by new
+ -- placeholder variables. For example:
+ -- C a (D (f x) (g y)) ==> C p1 (D p2 p3)
+
+argToPat :: ConstrEnv -> UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
+argToPat env us (Type ty)
+ = (us, Type ty)
+
+argToPat env us arg
+ | Just (CV dc args) <- is_con_app_maybe env arg
+ = let
+ (us',args') = argsToPats env us args
+ in
+ (us', mk_con_app dc args')
+
+argToPat env us (Var v) -- Don't uniqify existing vars,
+ = (us, Var v) -- so that we can spot when we pass them twice
+
+argToPat env us arg
+ = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
+ where
+ (us1,us2) = splitUniqSupply us
+
+argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
+argsToPats env us args = mapAccumL (argToPat env) us args
+\end{code}
+
+
+\begin{code}
+is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
+is_con_app_maybe env (Var v)
+ = lookupVarEnv env v
+ -- You might think we could look in the idUnfolding here
+ -- but that doesn't take account of which branch of a
+ -- case we are in, which is the whole point
+
+is_con_app_maybe env (Lit lit)
+ = Just (CV (LitAlt lit) [])
+
+is_con_app_maybe env expr
+ = case collectArgs expr of
+ (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
+ args `lengthAtLeast` dataConRepArity con
+ -- Might be > because the arity excludes type args
+ -> Just (CV (DataAlt con) args)
+
+ other -> Nothing
+
+mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
+mk_con_app (LitAlt lit) [] = Lit lit
+mk_con_app (DataAlt con) args = mkConApp con args
\end{code}
projects / ghc-hetmet.git / blobdiff