-specialise :: ScEnv
- -> Id -- Functionn
- -> [CoreBndr] -> CoreExpr -- Its RHS
- -> ScUsage -- Info on usage
- -> UniqSM ([CoreRule], -- Rules
- [(Id,CoreExpr)]) -- Bindings
-
-specialise env fn bndrs body body_usg
- = do { let (_, bndr_occs) = lookupOccs body_usg bndrs
-
- ; mb_calls <- mapM (callToPats (scope env) bndr_occs)
- (lookupVarEnv (calls body_usg) fn `orElse` [])
-
- ; let good_calls :: [([Var], [CoreArg])]
- good_calls = catMaybes mb_calls
- in_scope = mkInScopeSet $ unionVarSets $
- [ exprsFreeVars pats `delVarSetList` vs
- | (vs,pats) <- good_calls ]
- uniq_calls = nubBy (same_call in_scope) good_calls
- in
- mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
- (uniq_calls `zip` [1..]) }
- where
- -- Two calls are the same if they match both ways
- same_call 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
+type RhsInfo = (OutId, [OutVar], OutExpr, [ArgOcc])
+ -- Info about the *original* RHS of a binding we are specialising
+ -- Original binding f = \xs.body
+ -- Plus info about usage of arguments
+
+data SpecInfo = SI [OneSpec] -- The specialisations we have generated
+ Int -- Length of specs; used for numbering them
+ (Maybe ScUsage) -- Nothing => we have generated specialisations
+ -- from calls in the *original* RHS
+ -- Just cs => we haven't, and this is the usage
+ -- of the original RHS
+
+ -- One specialisation: Rule plus definition
+data OneSpec = OS CallPat -- Call pattern that generated this specialisation
+ CoreRule -- Rule connecting original id with the specialisation
+ OutId OutExpr -- Spec id + its rhs
+
+
+specLoop :: ScEnv
+ -> CallEnv
+ -> [RhsInfo]
+ -> ScUsage -> [SpecInfo] -- One per binder; acccumulating parameter
+ -> UniqSM (ScUsage, [SpecInfo]) -- ...ditto...
+specLoop env all_calls rhs_infos usg_so_far specs_so_far
+ = do { specs_w_usg <- zipWithM (specialise env all_calls) rhs_infos specs_so_far
+ ; let (new_usg_s, all_specs) = unzip specs_w_usg
+ new_usg = combineUsages new_usg_s
+ new_calls = scu_calls new_usg
+ all_usg = usg_so_far `combineUsage` new_usg
+ ; if isEmptyVarEnv new_calls then
+ return (all_usg, all_specs)
+ else
+ specLoop env new_calls rhs_infos all_usg all_specs }
+
+specialise
+ :: ScEnv
+ -> CallEnv -- Info on calls
+ -> RhsInfo
+ -> SpecInfo -- Original RHS plus patterns dealt with
+ -> UniqSM (ScUsage, SpecInfo) -- New specialised versions and their usage
+
+-- 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 bind_calls (fn, arg_bndrs, body, arg_occs)
+ spec_info@(SI specs spec_count mb_unspec)
+ | notNull arg_bndrs, -- Only specialise functions
+ Just all_calls <- lookupVarEnv bind_calls fn
+ = do { (boring_call, pats) <- callsToPats env specs arg_occs all_calls
+-- ; pprTrace "specialise" (vcat [ppr fn <+> ppr arg_occs,
+-- text "calls" <+> ppr all_calls,
+-- text "good pats" <+> ppr pats]) $
+-- return ()
+
+ -- Bale out if too many specialisations
+ -- Rather a hacky way to do so, but it'll do for now
+ ; let spec_count' = length pats + spec_count
+ ; case sc_count env of
+ Just max | spec_count' > max
+ -> pprTrace "SpecConstr: too many specialisations for one function (see -fspec-constr-count):"
+ (vcat [ptext (sLit "Function:") <+> ppr fn,
+ ptext (sLit "Specialisations:") <+> ppr (pats ++ [p | OS p _ _ _ <- specs])])
+ return (nullUsage, spec_info)
+
+ _normal_case -> do
+
+ { (spec_usgs, new_specs) <- mapAndUnzipM (spec_one env fn arg_bndrs body)
+ (pats `zip` [spec_count..])
+
+ ; let spec_usg = combineUsages spec_usgs
+ (new_usg, mb_unspec')
+ = case mb_unspec of
+ Just rhs_usg | boring_call -> (spec_usg `combineUsage` rhs_usg, Nothing)
+ _ -> (spec_usg, mb_unspec)
+
+ ; return (new_usg, SI (new_specs ++ specs) spec_count' mb_unspec') } }