import CoreSyn
import CoreLint ( showPass, endPass )
-import CoreUtils ( exprType, tcEqExpr, mkPiTypes )
+import CoreUtils ( exprType, mkPiTypes )
import CoreFVs ( exprsFreeVars )
-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 DataCon ( dataConRepArity, dataConUnivTyVars )
+import Type ( Type, tyConAppArgs )
+import Coercion ( coercionKind )
+import Rules ( matchN )
import Id ( Id, idName, idType, isDataConWorkId_maybe,
- mkUserLocal, mkSysLocal, idUnfolding )
+ mkUserLocal, mkSysLocal, idUnfolding, isLocalId )
import Var ( Var )
import VarEnv
import VarSet
import ErrUtils ( dumpIfSet_dyn )
import DynFlags ( DynFlags, DynFlag(..) )
import BasicTypes ( Activation(..) )
-import Maybes ( orElse )
-import Util ( mapAccumL, lengthAtLeast, notNull )
+import Maybes ( orElse, catMaybes, isJust )
+import Util ( zipWithEqual, lengthAtLeast, notNull )
import List ( nubBy, partition )
import UniqSupply
import Outputable
import FastString
+import UniqFM
\end{code}
-----------------------------------------------------
Hence the "OR" part of Note [Good arguments] below.
-ALTERNATIVE: pass both boxed and unboxed versions. This no longer saves
+ALTERNATIVE 2: pass both boxed and unboxed versions. This no longer saves
allocation, but does perhaps save evals. In the RULE we'd have
something like
rely on CSE to eliminate the duplicate allocation.... This alternative
doesn't look attractive enough to pursue.
+ALTERNATIVE 3: ignore the reboxing problem. The trouble is that
+the conservative reboxing story prevents many useful functions from being
+specialised. Example:
+ foo :: Maybe Int -> Int -> Int
+ foo (Just m) 0 = 0
+ foo x@(Just m) n = foo x (n-m)
+Here the use of 'x' will clearly not require boxing in the specialised function.
+
+The strictness analyser has the same problem, in fact. Example:
+ f p@(a,b) = ...
+If we pass just 'a' and 'b' to the worker, it might need to rebox the
+pair to create (a,b). A more sophisticated analysis might figure out
+precisely the cases in which this could happen, but the strictness
+analyser does no such analysis; it just passes 'a' and 'b', and hopes
+for the best.
+
+So my current choice is to make SpecConstr similarly aggressive, and
+ignore the bad potential of reboxing.
+
Note [Good arguments]
~~~~~~~~~~~~~~~~~~~~~
simplifier. That gives the simplest possible program for SpecConstr to
chew on; and it virtually guarantees no shadowing.
------------------------------------------------------
- Stuff not yet handled
------------------------------------------------------
-
-Here are notes arising from Roman's work that I don't want to lose.
-
-Specialising for constant parameters
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Note [Specialising for constant parameters]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This one is about specialising on a *constant* (but not necessarily
constructor) argument
Also
-Specialising for lambdas
-~~~~~~~~~~~~~~~~~~~~~~~~
+Note [Specialising for lambda parameters]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
foo :: Int -> (Int -> Int) -> Int
foo 0 f = 0
foo m f = foo (f m) (\n -> n-m)
Looks cool, but probably rare...but it might be easy to implement.
+
+Note [SpecConstr for casts]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider
+ data family T a :: *
+ data instance T Int = T Int
+
+ foo n = ...
+ where
+ go (T 0) = 0
+ go (T n) = go (T (n-1))
+
+The recursive call ends up looking like
+ go (T (I# ...) `cast` g)
+So we want to spot the construtor application inside the cast.
+That's why we have the Cast case in argToPat
+
+
+-----------------------------------------------------
+ Stuff not yet handled
+-----------------------------------------------------
+
+Here are notes arising from Roman's work that I don't want to lose.
+
Example 1
~~~~~~~~~
data T a = T !a
%************************************************************************
\begin{code}
-data ScEnv = SCE { scope :: VarEnv HowBound,
+data ScEnv = SCE { scope :: InScopeEnv,
-- Binds all non-top-level variables in scope
cons :: ConstrEnv
}
+type InScopeEnv = VarEnv HowBound
+
type ConstrEnv = IdEnv ConValue
data ConValue = CV AltCon [CoreArg]
-- Variables known to be bound to a constructor
instance Outputable ConValue where
ppr (CV con args) = ppr con <+> interpp'SP args
-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
+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
+ | RecArg -- These are those functions' arguments, or their sub-components;
+ -- we gather occurrence information for these
- | 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
+ | 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"
lookupScopeEnv env v = lookupVarEnv (scope env) v
-extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
+
+extendBndrsWith :: HowBound -> ScEnv -> [Var] -> ScEnv
+extendBndrsWith how_bound env bndrs
+ = env { scope = scope env `extendVarEnvList`
+ [(bndr,how_bound) | bndr <- bndrs] }
+
+extendBndrs env bndrs = extendBndrsWith Other env bndrs
extendBndr env bndr = env { scope = extendVarEnv (scope env) bndr Other }
-- When we encounter
-- 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
+extendCaseBndrs env case_bndr scrut con alt_bndrs
+ = case con of
+ DEFAULT -> env1
+ LitAlt lit -> extendCons env1 scrut case_bndr (CV con [])
+ DataAlt dc -> extendCons env1 scrut case_bndr (CV con vanilla_args)
+ where
+ vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
+ varsToCoreExprs alt_bndrs
where
- vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
- map varToCoreExpr alt_bndrs
-
- gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
- -- This call generates some bogus warnings from substExpr,
- -- because it's inconvenient to put all the Ids in scope
- -- Will be fixed when we move to FC
-
- (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]) }
+ env1 = extendBndrsWith (get_how scrut) env (case_bndr:alt_bndrs)
+
+ -- Record RecArg for the components iff the scrutinee is RecArg
+ -- I think the only reason for this is to keep the usage envt small
+ -- so is it worth it at all?
+ -- [This comment looks plain wrong to me, so I'm ignoring it
+ -- "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)... }" ]
+ get_how (Var v) = lookupVarEnv (scope env) v `orElse` Other
+ get_how (Cast e _) = get_how e
+ get_how (Note _ e) = get_how e
+ get_how other = Other
+
+extendCons :: ScEnv -> CoreExpr -> Id -> ConValue -> ScEnv
+extendCons env scrut case_bndr val
+ = case scrut of
+ Var v -> env { cons = extendVarEnv cons1 v val }
+ other -> env { cons = cons1 }
+ where
+ cons1 = extendVarEnv (cons env) case_bndr val
\end{code}
\begin{code}
data ScUsage
= SCU {
- calls :: !(IdEnv ([Call])), -- Calls
+ calls :: !(IdEnv [Call]), -- Calls
-- The functions are a subset of the
-- RecFuns in the ScEnv
combineUsages [] = nullUsage
combineUsages us = foldr1 combineUsage us
-data ArgOcc = CaseScrut
- | OtherOcc
- | Both
+lookupOcc :: ScUsage -> Var -> (ScUsage, ArgOcc)
+lookupOcc (SCU { calls = sc_calls, occs = sc_occs }) bndr
+ = (SCU {calls = sc_calls, occs = delVarEnv sc_occs bndr},
+ lookupVarEnv sc_occs bndr `orElse` NoOcc)
-instance Outputable ArgOcc where
- ppr CaseScrut = ptext SLIT("case-scrut")
- ppr OtherOcc = ptext SLIT("other-occ")
- ppr Both = ptext SLIT("case-scrut and other")
+lookupOccs :: ScUsage -> [Var] -> (ScUsage, [ArgOcc])
+lookupOccs (SCU { calls = sc_calls, occs = sc_occs }) bndrs
+ = (SCU {calls = sc_calls, occs = delVarEnvList sc_occs bndrs},
+ [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
+
+data ArgOcc = NoOcc -- Doesn't occur at all; or a type argument
+ | UnkOcc -- Used in some unknown way
+
+ | ScrutOcc (UniqFM [ArgOcc]) -- See Note [ScrutOcc]
+
+ | BothOcc -- Definitely taken apart, *and* perhaps used in some other way
+
+{- Note [ScrutOcc]
+
+An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
+is *only* taken apart or applied.
+
+ Functions, literal: ScrutOcc emptyUFM
+ Data constructors: ScrutOcc subs,
-combineOcc CaseScrut CaseScrut = CaseScrut
-combineOcc OtherOcc OtherOcc = OtherOcc
-combineOcc _ _ = Both
+where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
+The domain of the UniqFM is the Unique of the data constructor
+
+The [ArgOcc] is the occurrences of the *pattern-bound* components
+of the data structure. E.g.
+ data T a = forall b. MkT a b (b->a)
+A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
+
+-}
+
+instance Outputable ArgOcc where
+ ppr (ScrutOcc xs) = ptext SLIT("scrut-occ") <> ppr xs
+ ppr UnkOcc = ptext SLIT("unk-occ")
+ ppr BothOcc = ptext SLIT("both-occ")
+ ppr NoOcc = ptext SLIT("no-occ")
+
+-- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
+-- that if the thing is scrutinised anywhere then we get to see that
+-- in the overall result, even if it's also used in a boxed way
+-- This might be too agressive; see Note [Reboxing] Alternative 3
+combineOcc NoOcc occ = occ
+combineOcc occ NoOcc = occ
+combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
+combineOcc occ (ScrutOcc ys) = ScrutOcc ys
+combineOcc (ScrutOcc xs) occ = ScrutOcc xs
+combineOcc UnkOcc UnkOcc = UnkOcc
+combineOcc _ _ = BothOcc
+
+combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
+combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
+
+conArgOccs :: ArgOcc -> AltCon -> [ArgOcc]
+-- Find usage of components of data con; returns [UnkOcc...] if unknown
+-- See Note [ScrutOcc] for the extra UnkOccs in the vanilla datacon case
+
+conArgOccs (ScrutOcc fm) (DataAlt dc)
+ | Just pat_arg_occs <- lookupUFM fm dc
+ = [UnkOcc | tv <- dataConUnivTyVars dc] ++ pat_arg_occs
+
+conArgOccs other con = repeat UnkOcc
\end{code}
scExpr env e@(Type t) = returnUs (nullUsage, e)
scExpr env e@(Lit l) = returnUs (nullUsage, e)
-scExpr env e@(Var v) = returnUs (varUsage env v OtherOcc, e)
+scExpr env e@(Var v) = returnUs (varUsage env v UnkOcc, e)
scExpr env (Note n e) = scExpr env e `thenUs` \ (usg,e') ->
returnUs (usg, Note n e')
+scExpr env (Cast e co)= scExpr env e `thenUs` \ (usg,e') ->
+ returnUs (usg, Cast e' co)
scExpr env (Lam b e) = scExpr (extendBndr env b) e `thenUs` \ (usg,e') ->
returnUs (usg, Lam b e')
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 ty alts')
+ = do { (alt_usgs, alt_occs, alts') <- mapAndUnzip3Us sc_alt alts
+ ; let (alt_usg, b_occ) = lookupOcc (combineUsages alt_usgs) b
+ scrut_occ = foldr combineOcc b_occ alt_occs
+ -- The combined usage of the scrutinee is given
+ -- by scrut_occ, which is passed to scScrut, which
+ -- in turn treats a bare-variable scrutinee specially
+ ; (scrut_usg, scrut') <- scScrut env scrut scrut_occ
+ ; return (alt_usg `combineUsage` scrut_usg,
+ 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 env1 rhs `thenUs` \ (usg,rhs') ->
- returnUs (usg, (con,bs,rhs'))
- where
- env1 = extendCaseBndrs env b scrut con bs
+ sc_alt (con,bs,rhs)
+ = do { let env1 = extendCaseBndrs env b scrut con bs
+ ; (usg,rhs') <- scExpr env1 rhs
+ ; let (usg', arg_occs) = lookupOccs usg bs
+ scrut_occ = case con of
+ DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
+ other -> ScrutOcc emptyUFM
+ ; return (usg', scrut_occ, (con,bs,rhs')) }
scExpr env (Let bind body)
= scBind env bind `thenUs` \ (env', bind_usg, bind') ->
returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
scExpr env e@(App _ _)
- = let
- (fn, args) = collectArgs e
- in
- mapAndUnzipUs (scExpr env) (fn:args) `thenUs` \ (usgs, (fn':args')) ->
+ = do { let (fn, args) = collectArgs e
+ ; (fn_usg, fn') <- scScrut env fn (ScrutOcc emptyUFM)
-- Process the function too. It's almost always a variable,
-- but not always. In particular, if this pass follows float-in,
-- which it may, we can get
-- (let f = ...f... in f) arg1 arg2
- let
- call_usg = case fn of
- Var f | Just RecFun <- lookupScopeEnv env f
- -> SCU { calls = unitVarEnv f [(cons env, args)],
- occs = emptyVarEnv }
- other -> nullUsage
- in
- returnUs (combineUsages usgs `combineUsage` call_usg, mkApps fn' args')
+ -- We use scScrut to record the fact that the function is called
+ -- Perhpas we should check that it has at least one value arg,
+ -- but currently we don't bother
+
+ ; (arg_usgs, args') <- mapAndUnzipUs (scExpr env) args
+ ; let call_usg = case fn of
+ Var f | Just RecFun <- lookupScopeEnv env f
+ , not (null args) -- Not a proper call!
+ -> SCU { calls = unitVarEnv f [(cons env, args)],
+ occs = emptyVarEnv }
+ other -> nullUsage
+ ; return (combineUsages arg_usgs `combineUsage` fn_usg
+ `combineUsage` call_usg,
+ mkApps fn' args') }
----------------------
-scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
-scBind env (Rec [(fn,rhs)])
- | notNull val_bndrs
- = scExpr env_fn_body body `thenUs` \ (usg, body') ->
- specialise env fn bndrs body' usg `thenUs` \ (rules, spec_prs) ->
- -- Note body': the specialised copies should be based on the
- -- optimised version of the body, in case there were
- -- nested functions inside.
- let
- SCU { calls = calls, occs = occs } = usg
- in
- 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_fn_body = extendRecBndr env fn bndrs
+scScrut :: ScEnv -> CoreExpr -> ArgOcc -> UniqSM (ScUsage, CoreExpr)
+-- Used for the scrutinee of a case,
+-- or the function of an application.
+-- Remember to look through casts
+scScrut env e@(Var v) occ = returnUs (varUsage env v occ, e)
+scScrut env (Cast e co) occ = do { (usg, e') <- scScrut env e occ
+ ; returnUs (usg, Cast e' co) }
+scScrut env e occ = scExpr env e
+
+----------------------
+scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
scBind env (Rec prs)
- = mapAndUnzipUs do_one prs `thenUs` \ (usgs, prs') ->
- returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
- where
- do_one (bndr,rhs) = scExpr env rhs `thenUs` \ (usg, rhs') ->
- returnUs (usg, (bndr,rhs'))
+ = do { let bndrs = map fst prs
+ rhs_env = extendBndrsWith RecFun env bndrs
+
+ ; (rhs_usgs, prs_w_occs) <- mapAndUnzipUs (scRecRhs rhs_env) prs
+ ; let rhs_usg = combineUsages rhs_usgs
+ rhs_calls = calls rhs_usg
+
+ ; prs_s <- mapUs (specialise env rhs_calls) prs_w_occs
+ ; return (extendBndrs env bndrs,
+ -- For the body of the letrec, just
+ -- extend the env with Other to record
+ -- that it's in scope; no funny RecFun business
+ rhs_usg { calls = calls rhs_usg `delVarEnvList` bndrs },
+ Rec (concat prs_s)) }
scBind env (NonRec bndr rhs)
- = scExpr env rhs `thenUs` \ (usg, rhs') ->
- returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
+ = do { (usg, rhs') <- scExpr env rhs
+ ; return (extendBndr env bndr, usg, NonRec bndr rhs') }
+
+----------------------
+scRecRhs :: ScEnv -> (Id,CoreExpr)
+ -> UniqSM (ScUsage, (Id, CoreExpr, [ArgOcc]))
+-- The returned [ArgOcc] says how the visible,
+-- lambda-bound binders of the RHS are used
+-- (including the TyVar binders)
+scRecRhs env (bndr,rhs)
+ = do { let (arg_bndrs,body) = collectBinders rhs
+ body_env = extendBndrsWith RecArg env arg_bndrs
+ ; (body_usg, body') <- scExpr body_env body
+ ; let (rhs_usg, arg_occs) = lookupOccs body_usg arg_bndrs
+ ; return (rhs_usg, (bndr, mkLams arg_bndrs body', arg_occs)) }
----------------------
varUsage env v use
%************************************************************************
\begin{code}
-specialise :: ScEnv
- -> Id -- Functionn
- -> [CoreBndr] -> CoreExpr -- Its RHS
- -> ScUsage -- Info on usage
- -> UniqSM ([CoreRule], -- Rules
- [(Id,CoreExpr)]) -- Bindings
-
-specialise env fn bndrs body (SCU {calls=calls, occs=occs})
- = getUs `thenUs` \ us ->
- let
- all_calls = lookupVarEnv calls fn `orElse` []
-
- good_calls :: [[CoreArg]]
- good_calls = [ pats
- | (con_env, call_args) <- all_calls,
- call_args `lengthAtLeast` n_bndrs, -- App is saturated
- let call = bndrs `zip` 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
- mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
- (nubBy same_call good_calls `zip` [1..])
+specialise
+ :: ScEnv
+ -> IdEnv [Call] -- Info on usage
+ -> (Id, CoreExpr, [ArgOcc]) -- Original binding, plus info on how the rhs's
+ -- lambda-binders are used (includes TyVar bndrs)
+ -> UniqSM [(Id,CoreExpr)] -- Original binding (decorated with rules)
+ -- plus specialised bindings
+
+-- Note: the rhs here is the optimised version of the original rhs
+-- So when we make a specialised copy of the RHS, we're starting
+-- from an RHS whose nested functions have been optimised already.
+
+specialise env calls (fn, rhs, arg_occs)
+ | notNull arg_occs, -- Only specialise functions
+ Just all_calls <- lookupVarEnv calls fn
+ = do { mb_pats <- mapM (callToPats (scope env) arg_occs) all_calls
+
+ ; let good_pats :: [([Var], [CoreArg])]
+ good_pats = catMaybes mb_pats
+ in_scope = mkInScopeSet $ unionVarSets $
+ [ exprsFreeVars pats `delVarSetList` vs
+ | (vs,pats) <- good_pats ]
+ uniq_pats = nubBy (same_pat in_scope) good_pats
+-- ; pprTrace "specialise" (vcat [ppr fn <+> ppr arg_occs,
+-- text "calls" <+> ppr all_calls,
+-- text "good pats" <+> ppr good_pats,
+-- text "uniq pats" <+> ppr uniq_pats]) $
+-- return ()
+
+ ; (rules, spec_prs) <- mapAndUnzipUs (spec_one fn rhs)
+ (uniq_pats `zip` [1..])
+
+ ; return ((fn `addIdSpecialisations` rules, rhs) : spec_prs) }
+
+ | otherwise
+ = return [(fn,rhs)] -- The boring case
where
- n_bndrs = length bndrs
- same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
-
----------------------
-good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
--- See Note [Good arguments] above
-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
- = 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
- other -> False
- other -> False -- Not used, or used wonkily
-
+ -- Two pats are the same if they match both ways
+ same_pat in_scope (vs1,as1)(vs2,as2)
+ = isJust (matchN in_scope vs1 as1 as2)
+ && isJust (matchN in_scope vs2 as2 as1)
+
+callToPats :: InScopeEnv -> [ArgOcc] -> Call
+ -> UniqSM (Maybe ([Var], [CoreExpr]))
+ -- The VarSet is the variables to quantify over in the rule
+ -- The [CoreExpr] are the argument patterns for the rule
+callToPats in_scope bndr_occs (con_env, args)
+ | length args < length bndr_occs -- Check saturated
+ = return Nothing
+ | otherwise
+ = do { prs <- argsToPats in_scope con_env (args `zip` bndr_occs)
+ ; let (good_pats, pats) = unzip prs
+ pat_fvs = varSetElems (exprsFreeVars pats)
+ qvars = filter (not . (`elemVarEnv` in_scope)) pat_fvs
+ -- Quantify over variables that are not in sccpe
+ -- See Note [Shadowing] at the top
+
+ ; -- pprTrace "callToPats" (ppr args $$ ppr prs $$ ppr bndr_occs) $
+ if or good_pats
+ then return (Just (qvars, pats))
+ else return Nothing }
---------------------
-spec_one :: ScEnv
- -> Id -- Function
+spec_one :: Id -- Function
-> CoreExpr -- Rhs of the original function
- -> ([CoreArg], Int)
+ -> (([Var], [CoreArg]), Int)
-> UniqSM (CoreRule, (Id,CoreExpr)) -- Rule and binding
-- spec_one creates a specialised copy of the function, together
f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
-}
-spec_one env fn rhs (pats, rule_number)
+spec_one fn rhs ((vars_to_bind, 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
- -- See Note [Shadowing] at the top
- 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
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
-- 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)))
+argToPat :: InScopeEnv -- What's in scope at the fn defn site
+ -> ConstrEnv -- ConstrEnv at the call site
+ -> CoreArg -- A call arg (or component thereof)
+ -> ArgOcc
+ -> UniqSM (Bool, CoreArg)
+-- Returns (interesting, pat),
+-- where pat is the pattern derived from the argument
+-- intersting=True if the pattern is non-trivial (not a variable or type)
+-- E.g. x:xs --> (True, x:xs)
+-- f xs --> (False, w) where w is a fresh wildcard
+-- (f xs, 'c') --> (True, (w, 'c')) where w is a fresh wildcard
+-- \x. x+y --> (True, \x. x+y)
+-- lvl7 --> (True, lvl7) if lvl7 is bound
+-- somewhere further out
+
+argToPat in_scope con_env arg@(Type ty) arg_occ
+ = return (False, arg)
+
+argToPat in_scope con_env (Let _ arg) arg_occ
+ = argToPat in_scope con_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
+ ; if interesting then
+ return (interesting, Cast arg' co)
+ else
+ wildCardPat (snd (coercionKind co)) }
+
+argToPat in_scope con_env arg arg_occ
+ | is_value_lam arg
+ = return (True, arg)
where
- (us1,us2) = splitUniqSupply us
-
-argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
-argsToPats env us args = mapAccumL (argToPat env) us args
+ is_value_lam (Lam v e) -- Spot a value lambda, even if
+ | isId v = True -- it is inside a type lambda
+ | otherwise = is_value_lam e
+ is_value_lam other = False
+
+ -- 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) <- is_con_app_maybe con_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)
+ ; return (True, mk_con_app dc (map snd args')) }
+
+ -- Check if the argument is a variable that
+ -- is in scope at the function definition site
+ -- 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 `elemVarEnv` in_scope,
+ case arg_occ of { UnkOcc -> False; other -> True }, -- (a)
+ isValueUnfolding (idUnfolding v) -- (b)
+ = return (True, Var v)
+
+ -- Check for a variable bound inside the function.
+ -- 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
+ = return (False, Var v)
+
+ -- The default case: make a wild-card
+argToPat in_scope con_env arg arg_occ
+ = wildCardPat (exprType arg)
+
+wildCardPat :: Type -> UniqSM (Bool, CoreArg)
+wildCardPat ty = do { uniq <- getUniqueUs
+ ; let id = mkSysLocal FSLIT("sc") uniq ty
+ ; return (False, Var id) }
+
+argsToPats :: InScopeEnv -> ConstrEnv
+ -> [(CoreArg, ArgOcc)]
+ -> UniqSM [(Bool, CoreArg)]
+argsToPats in_scope con_env args
+ = mapUs do_one args
+ where
+ do_one (arg,occ) = argToPat in_scope con_env arg occ
\end{code}
\begin{code}
is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
+is_con_app_maybe env (Lit lit)
+ = Just (CV (LitAlt lit) [])
+
+is_con_app_maybe 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)
+
is_con_app_maybe env (Var v)
- = case lookupVarEnv env v of
- Just stuff -> Just stuff
- -- You might think we could look in the idUnfolding here
+ | 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
- Nothing | isCheapUnfolding unf
- -> is_con_app_maybe env (unfoldingTemplate unf)
- where
- unf = idUnfolding v
- -- However we do want to consult the unfolding as well,
- -- for let-bound constructors!
-
- other -> Nothing
-
-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)
+ | isCheapUnfolding unf
+ = is_con_app_maybe env (unfoldingTemplate unf)
+ where
+ unf = idUnfolding v
+ -- However we do want to consult the unfolding
+ -- as well, for let-bound constructors!
- other -> Nothing
+is_con_app_maybe env expr = Nothing
mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
mk_con_app (LitAlt lit) [] = Lit lit
mk_con_app (DataAlt con) args = mkConApp con args
+mk_con_app other args = panic "SpecConstr.mk_con_app"
\end{code}