import CoreFVs -- all of it
import CoreSubst ( Subst, emptySubst, extendInScope, extendInScopeList,
extendIdSubst, cloneIdBndr, cloneRecIdBndrs )
-import Id ( idType, mkLocalIdWithInfo, mkSysLocal, isOneShotLambda,
- zapDemandIdInfo, transferPolyIdInfo,
- idSpecialisation, idUnfolding, setIdInfo,
- setIdStrictness, setIdArity
- )
+import Id
import IdInfo
import Var
import VarSet
lvlExpr _ env (_, AnnVar v) = return (lookupVar env v)
lvlExpr _ _ (_, AnnLit lit) = return (Lit lit)
-lvlExpr ctxt_lvl env (_, AnnApp fun arg) = do
- fun' <- lvlExpr ctxt_lvl env fun -- We don't do MFE on partial applications
- arg' <- lvlMFE False ctxt_lvl env arg
- return (App fun' arg')
+lvlExpr ctxt_lvl env expr@(_, AnnApp _ _) = do
+ let
+ (fun, args) = collectAnnArgs expr
+ --
+ case fun of
+ -- float out partial applications. This is very beneficial
+ -- in some cases (-7% runtime -4% alloc over nofib -O2).
+ -- In order to float a PAP, there must be a function at the
+ -- head of the application, and the application must be
+ -- over-saturated with respect to the function's arity.
+ (_, AnnVar f) | floatPAPs env &&
+ arity > 0 && arity < n_val_args ->
+ do
+ let (lapp, rargs) = left (n_val_args - arity) expr []
+ rargs' <- mapM (lvlMFE False ctxt_lvl env) rargs
+ lapp' <- lvlMFE False ctxt_lvl env lapp
+ return (foldl App lapp' rargs')
+ where
+ n_val_args = count (isValArg . deAnnotate) args
+ arity = idArity f
+
+ -- separate out the PAP that we are floating from the extra
+ -- arguments, by traversing the spine until we have collected
+ -- (n_val_args - arity) value arguments.
+ left 0 e rargs = (e, rargs)
+ left n (_, AnnApp f a) rargs
+ | isValArg (deAnnotate a) = left (n-1) f (a:rargs)
+ | otherwise = left n f (a:rargs)
+ left _ _ _ = panic "SetLevels.lvlExpr.left"
+
+ -- No PAPs that we can float: just carry on with the
+ -- arguments and the function.
+ _otherwise -> do
+ args' <- mapM (lvlMFE False ctxt_lvl env) args
+ fun' <- lvlExpr ctxt_lvl env fun
+ return (foldl App fun' args')
lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
expr' <- lvlExpr ctxt_lvl env expr
-> LvlM (LevelledBind, LevelEnv)
lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
- | isTyVar bndr -- Don't do anything for TyVar binders
+ | isTyCoVar bndr -- Don't do anything for TyVar binders
-- (simplifier gets rid of them pronto)
= do rhs' <- lvlExpr ctxt_lvl env rhs
return (NonRec (TB bndr ctxt_lvl) rhs', env)
= (lvl, [])
lvlLamBndrs lvl bndrs
- = go (incMinorLvl lvl)
- False -- Havn't bumped major level in this group
- [] bndrs
+ = (new_lvl, [TB bndr new_lvl | bndr <- bndrs])
+ -- All the new binders get the same level, because
+ -- any floating binding is either going to float past
+ -- all or none. We never separate binders
where
- go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
- | isId bndr && -- Go to the next major level if this is a value binder,
- not bumped_major && -- and we havn't already gone to the next level (one jump per group)
- not (isOneShotLambda bndr) -- and it isn't a one-shot lambda
- = go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
-
- | otherwise
- = go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
+ new_lvl | any is_major bndrs = incMajorLvl lvl
+ | otherwise = incMinorLvl lvl
- where
- new_lvl = incMajorLvl old_lvl
-
- go old_lvl _ rev_lvld_bndrs []
- = (old_lvl, reverse rev_lvld_bndrs)
- -- a lambda like this (\x -> coerce t (\s -> ...))
- -- This happens quite a bit in state-transformer programs
+ is_major bndr = isId bndr && not (isOneShotLambda bndr)
\end{code}
\begin{code}
floatConsts :: LevelEnv -> Bool
floatConsts (fos, _, _, _) = floatOutConstants fos
+floatPAPs :: LevelEnv -> Bool
+floatPAPs (fos, _, _, _) = floatOutPartialApplications fos
+
extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
-- Used when *not* cloning
extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
(False, True) -> False
_ -> v1 <= v2 -- Same family
- is_tv v = isTyVar v && not (isCoVar v)
+ is_tv v = isTyCoVar v && not (isCoVar v)
uniq :: [Var] -> [Var]
-- Remove adjacent duplicates; the sort will have brought them together
-- We are going to lambda-abstract, so nuke any IdInfo,
-- and add the tyvars of the Id (if necessary)
- zap v | isId v = WARN( isInlineRule (idUnfolding v) ||
+ zap v | isId v = WARN( isStableUnfolding (idUnfolding v) ||
not (isEmptySpecInfo (idSpecialisation v)),
text "absVarsOf: discarding info on" <+> ppr v )
setIdInfo v vanillaIdInfo