From 57a4597d5a487af65d0a6c9a6701e2efcbbefac2 Mon Sep 17 00:00:00 2001 From: "simonpj@microsoft.com" Date: Sun, 5 Aug 2007 22:03:35 +0000 Subject: [PATCH] Make SpecConstr specialise for constant arguments again Consider lvl = Just True foo :: Maybe Bool -> Int -> Int foo (Just True) i = i foo _ i = foo lvl i SpecConstr should specialise foo, but it wasn't doing so (spotted by Roman). Reason: lvl's unfolding wasn't in the cloned version of lvl. Solution: extend the value environment to record top-level bindings too At the same time I made it work if 'lvl' is a lambda, in which case it is again worth specialisg. This meant renaming ConEnv to ValueEnv, and adding a case for 'LambdaVal'. (To make specialisation on lambdas work properly, we have to do lambda lifting as well, but this gets part of the way, and fixes a bug too.) --- compiler/specialise/SpecConstr.lhs | 148 ++++++++++++++++++++---------------- 1 file changed, 83 insertions(+), 65 deletions(-) diff --git a/compiler/specialise/SpecConstr.lhs b/compiler/specialise/SpecConstr.lhs index db06d55..48718af 100644 --- a/compiler/specialise/SpecConstr.lhs +++ b/compiler/specialise/SpecConstr.lhs @@ -22,7 +22,7 @@ import WwLib ( mkWorkerArgs ) import DataCon ( dataConRepArity, dataConUnivTyVars ) import Type ( Type, tyConAppArgs ) import Coercion ( coercionKind ) -import Id ( Id, idName, idType, isDataConWorkId_maybe, +import Id ( Id, idName, idType, isDataConWorkId_maybe, idArity, mkUserLocal, mkSysLocal, idUnfolding, isLocalId ) import Var ( Var ) import VarEnv @@ -33,7 +33,7 @@ import OccName ( mkSpecOcc ) import ErrUtils ( dumpIfSet_dyn ) import DynFlags ( DynFlags(..), DynFlag(..) ) import BasicTypes ( Activation(..) ) -import Maybes ( orElse, catMaybes ) +import Maybes ( orElse, catMaybes, isJust ) import Util import List ( nubBy, partition ) import UniqSupply @@ -457,8 +457,9 @@ data ScEnv = SCE { sc_size :: Int, -- Size threshold -- Binds interesting non-top-level variables -- Domain is OutVars (*after* applying the substitution) - sc_cons :: ConstrEnv + sc_vals :: ValueEnv -- Domain is OutIds (*after* applying the substitution) + -- Used even for top-level bindings (but not imported ones) } --------------------- @@ -473,21 +474,20 @@ type OutVar = Var type HowBoundEnv = VarEnv HowBound -- Domain is OutVars --------------------- -type ConstrEnv = IdEnv ConValue -- Domain is OutIds -data ConValue = CV AltCon [CoreArg] - -- Variables known to be bound to a constructor - -- in a particular case alternative +type ValueEnv = IdEnv Value -- Domain is OutIds +data Value = ConVal AltCon [CoreArg] -- *Saturated* constructors + | LambdaVal -- Inlinable lambdas or PAPs - -instance Outputable ConValue where - ppr (CV con args) = ppr con <+> interpp'SP args +instance Outputable Value where + ppr (ConVal con args) = ppr con <+> interpp'SP args + ppr LambdaVal = ptext SLIT("") --------------------- initScEnv dflags = SCE { sc_size = specThreshold dflags, sc_subst = emptySubst, sc_how_bound = emptyVarEnv, - sc_cons = emptyVarEnv } + sc_vals = emptyVarEnv } data HowBound = RecFun -- These are the recursive functions for which -- we seek interesting call patterns @@ -549,26 +549,26 @@ extendBndr env bndr = (env { sc_subst = subst' }, bndr') where (subst', bndr') = substBndr (sc_subst env) bndr -extendConEnv :: ScEnv -> Id -> Maybe ConValue -> ScEnv -extendConEnv env id Nothing = env -extendConEnv env id (Just cv) = env { sc_cons = extendVarEnv (sc_cons env) id cv } +extendValEnv :: ScEnv -> Id -> Maybe Value -> ScEnv +extendValEnv env id Nothing = env +extendValEnv env id (Just cv) = env { sc_vals = extendVarEnv (sc_vals env) id cv } extendCaseBndrs :: ScEnv -> CoreExpr -> Id -> AltCon -> [Var] -> ScEnv -- When we encounter -- case scrut of b -- C x y -> ... -- we want to bind b, and perhaps scrut too, to (C x y) --- NB: Extends only the sc_cons part of the envt +-- NB: Extends only the sc_vals part of the envt extendCaseBndrs env scrut case_bndr con alt_bndrs = case scrut of - Var v -> extendConEnv env1 v cval + Var v -> extendValEnv env1 v cval other -> env1 where - env1 = extendConEnv env case_bndr cval + env1 = extendValEnv env case_bndr cval cval = case con of DEFAULT -> Nothing - LitAlt lit -> Just (CV con []) - DataAlt dc -> Just (CV con vanilla_args) + LitAlt lit -> Just (ConVal con []) + DataAlt dc -> Just (ConVal con vanilla_args) where vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++ varsToCoreExprs alt_bndrs @@ -593,7 +593,7 @@ data ScUsage -- RecArg in the ScEnv type CallEnv = IdEnv [Call] -type Call = (ConstrEnv, [CoreArg]) +type Call = (ValueEnv, [CoreArg]) -- The arguments of the call, together with the -- env giving the constructor bindings at the call site @@ -719,16 +719,15 @@ scExpr' env (Lam b e) = do { let (env', b') = extendBndr env b scExpr' env (Case scrut b ty alts) = do { (scrut_usg, scrut') <- scExpr env scrut - ; case isConApp (sc_cons env) scrut' of - Nothing -> sc_vanilla scrut_usg scrut' - Just cval -> sc_con_app cval scrut' + ; case isValue (sc_vals env) scrut' of + Just (ConVal con args) -> sc_con_app con args scrut' + other -> sc_vanilla scrut_usg scrut' } where - sc_con_app cval@(CV con args) scrut' -- Known constructor; simplify + sc_con_app con args scrut' -- Known constructor; simplify = do { let (_, bs, rhs) = findAlt con alts alt_env' = extendScSubst env ((b,scrut') : bs `zip` trimConArgs con args) ; scExpr alt_env' rhs } - sc_vanilla scrut_usg scrut' -- Normal case = do { let (alt_env,b') = extendBndrWith RecArg env b @@ -764,8 +763,8 @@ scExpr' env (Let (NonRec bndr rhs) body) ; if null args' || isEmptyVarEnv (calls rhs_usg) then do do { -- Vanilla case let rhs' = mkLams args' rhs_body' - body_env2 = extendConEnv body_env bndr' (isConApp (sc_cons env) rhs') - -- Record if the RHS is a constructor + body_env2 = extendValEnv body_env bndr' (isValue (sc_vals env) rhs') + -- Record if the RHS is a value ; (body_usg, body') <- scExpr body_env2 body ; return (body_usg `combineUsage` rhs_usg, Let (NonRec bndr' rhs') body') } else @@ -807,7 +806,7 @@ scExpr' env e@(App _ _) ; let call_usg = case fn' of Var f | Just RecFun <- lookupHowBound env f , not (null args) -- Not a proper call! - -> SCU { calls = unitVarEnv f [(sc_cons env, args')], + -> SCU { calls = unitVarEnv f [(sc_vals env, args')], occs = emptyVarEnv } other -> nullUsage ; return (combineUsages arg_usgs `combineUsage` fn_usg' @@ -859,8 +858,9 @@ scBind env (Rec prs) scBind env (NonRec bndr rhs) = do { (usg, rhs') <- scExpr env rhs - ; let (env', bndr') = extendBndr env bndr - ; return (env', usg, NonRec bndr' rhs') } + ; let (env1, bndr') = extendBndr env bndr + env2 = extendValEnv env1 bndr' (isValue (sc_vals env) rhs') + ; return (env2, usg, NonRec bndr' rhs') } ---------------------- scRecRhs :: ScEnv -> (OutId, InExpr) -> UniqSM (ScUsage, RhsInfo) @@ -1065,7 +1065,7 @@ callToPats env bndr_occs (con_env, args) -- C a (D (f x) (g y)) ==> C p1 (D p2 p3) argToPat :: InScopeSet -- What's in scope at the fn defn site - -> ConstrEnv -- ConstrEnv at the call site + -> ValueEnv -- ValueEnv at the call site -> CoreArg -- A call arg (or component thereof) -> ArgOcc -> UniqSM (Bool, CoreArg) @@ -1079,11 +1079,11 @@ argToPat :: InScopeSet -- What's in scope at the fn defn site -- lvl7 --> (True, lvl7) if lvl7 is bound -- somewhere further out -argToPat in_scope con_env arg@(Type ty) arg_occ +argToPat in_scope val_env arg@(Type ty) arg_occ = return (False, arg) -argToPat in_scope con_env (Note n arg) arg_occ - = argToPat in_scope con_env arg arg_occ +argToPat in_scope val_env (Note n arg) arg_occ + = argToPat in_scope val_env arg arg_occ -- Note [Notes in call patterns] -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -- Ignore Notes. In particular, we want to ignore any InlineMe notes @@ -1091,15 +1091,15 @@ argToPat in_scope con_env (Note n arg) arg_occ -- ride roughshod over them all for now. --- See Note [Notes in RULE matching] in Rules -argToPat in_scope con_env (Let _ arg) arg_occ - = argToPat in_scope con_env arg arg_occ +argToPat in_scope val_env (Let _ arg) arg_occ + = argToPat in_scope val_env arg arg_occ -- Look through let expressions -- e.g. f (let v = rhs in \y -> ...v...) -- Here we can specialise for f (\y -> ...) -- because the rule-matcher will look through the let. -argToPat in_scope con_env (Cast arg co) arg_occ - = do { (interesting, arg') <- argToPat in_scope con_env arg arg_occ +argToPat in_scope val_env (Cast arg co) arg_occ + = do { (interesting, arg') <- argToPat in_scope val_env arg arg_occ ; if interesting then return (interesting, Cast arg' co) else @@ -1107,7 +1107,7 @@ argToPat in_scope con_env (Cast arg co) arg_occ {- Disabling lambda specialisation for now It's fragile, and the spec_loop can be infinite -argToPat in_scope con_env arg arg_occ +argToPat in_scope val_env arg arg_occ | is_value_lam arg = return (True, arg) where @@ -1119,15 +1119,15 @@ argToPat in_scope con_env arg arg_occ -- Check for a constructor application -- NB: this *precedes* the Var case, so that we catch nullary constrs -argToPat in_scope con_env arg arg_occ - | Just (CV dc args) <- isConApp con_env arg +argToPat in_scope val_env arg arg_occ + | Just (ConVal dc args) <- isValue val_env arg , case arg_occ of ScrutOcc _ -> True -- Used only by case scrutinee BothOcc -> case arg of -- Used elsewhere App {} -> True -- see Note [Reboxing] other -> False other -> False -- No point; the arg is not decomposed - = do { args' <- argsToPats in_scope con_env (args `zip` conArgOccs arg_occ dc) + = do { args' <- argsToPats in_scope val_env (args `zip` conArgOccs arg_occ dc) ; return (True, mk_con_app dc (map snd args')) } -- Check if the argument is a variable that @@ -1135,11 +1135,17 @@ argToPat in_scope con_env arg arg_occ -- It's worth specialising on this if -- (a) it's used in an interesting way in the body -- (b) we know what its value is -argToPat in_scope con_env (Var v) arg_occ - | not (isLocalId v) || v `elemInScopeSet` in_scope, - case arg_occ of { UnkOcc -> False; other -> True }, -- (a) - isValueUnfolding (idUnfolding v) -- (b) +argToPat in_scope val_env (Var v) arg_occ + | case arg_occ of { UnkOcc -> False; other -> True }, -- (a) + is_value -- (b) = return (True, Var v) + where + is_value + | isLocalId v = v `elemInScopeSet` in_scope + && isJust (lookupVarEnv val_env v) + -- Local variables have values in val_env + | otherwise = isValueUnfolding (idUnfolding v) + -- Imports have unfoldings -- I'm really not sure what this comment means -- And by not wild-carding we tend to get forall'd @@ -1150,11 +1156,11 @@ argToPat in_scope con_env (Var v) arg_occ -- Don't make a wild-card, because we may usefully share -- e.g. f a = let x = ... in f (x,x) -- NB: this case follows the lambda and con-app cases!! -argToPat in_scope con_env (Var v) arg_occ +argToPat in_scope val_env (Var v) arg_occ = return (False, Var v) -- The default case: make a wild-card -argToPat in_scope con_env arg arg_occ +argToPat in_scope val_env arg arg_occ = wildCardPat (exprType arg) wildCardPat :: Type -> UniqSM (Bool, CoreArg) @@ -1162,42 +1168,54 @@ wildCardPat ty = do { uniq <- getUniqueUs ; let id = mkSysLocal FSLIT("sc") uniq ty ; return (False, Var id) } -argsToPats :: InScopeSet -> ConstrEnv +argsToPats :: InScopeSet -> ValueEnv -> [(CoreArg, ArgOcc)] -> UniqSM [(Bool, CoreArg)] -argsToPats in_scope con_env args +argsToPats in_scope val_env args = mapUs do_one args where - do_one (arg,occ) = argToPat in_scope con_env arg occ + do_one (arg,occ) = argToPat in_scope val_env arg occ \end{code} \begin{code} -isConApp :: ConstrEnv -> CoreExpr -> Maybe ConValue -isConApp env (Lit lit) - = Just (CV (LitAlt lit) []) - -isConApp env expr -- Maybe it's a constructor application - | (Var fun, args) <- collectArgs expr, - Just con <- isDataConWorkId_maybe fun, - args `lengthAtLeast` dataConRepArity con - -- Might be > because the arity excludes type args - = Just (CV (DataAlt con) args) - -isConApp env (Var v) +isValue :: ValueEnv -> CoreExpr -> Maybe Value +isValue env (Lit lit) + = Just (ConVal (LitAlt lit) []) + +isValue env (Var v) | Just stuff <- lookupVarEnv env v = Just stuff -- 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 | not (isLocalId v) && isCheapUnfolding unf - = isConApp env (unfoldingTemplate unf) + = isValue env (unfoldingTemplate unf) where unf = idUnfolding v -- However we do want to consult the unfolding -- as well, for let-bound constructors! -isConApp env expr = Nothing +isValue env (Lam b e) + | isTyVar b = isValue env e + | otherwise = Just LambdaVal + +isValue env expr -- Maybe it's a constructor application + | (Var fun, args) <- collectArgs expr + = case isDataConWorkId_maybe fun of + + Just con | args `lengthAtLeast` dataConRepArity con + -- Check saturated; might be > because the + -- arity excludes type args + -> Just (ConVal (DataAlt con) args) + + other | valArgCount args < idArity fun + -- Under-applied function + -> Just LambdaVal -- Partial application + + other -> Nothing + +isValue env expr = Nothing mk_con_app :: AltCon -> [CoreArg] -> CoreExpr mk_con_app (LitAlt lit) [] = Lit lit -- 1.7.10.4