ghc/compiler/stranal/WorkWrap.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
   3 %
   4 \section[WorkWrap]{Worker/wrapper-generating back-end of strictness analyser}
   5
   6 \begin{code}
   7 module WorkWrap ( wwTopBinds, getWorkerIdAndCons ) where
   8
   9 #include "HsVersions.h"
  10
  11 import CoreSyn
  12 import CoreUnfold       ( Unfolding, certainlySmallEnoughToInline, calcUnfoldingGuidance )
  13 import CmdLineOpts      ( opt_UnfoldingCreationThreshold, opt_D_verbose_core2core,
  14                           opt_D_dump_worker_wrapper )
  15 import CoreLint         ( beginPass, endPass )
  16 import CoreUtils        ( coreExprType )
  17 import Const            ( Con(..) )
  18 import DataCon          ( DataCon )
  19 import MkId             ( mkWorkerId )
  20 import Id               ( Id, getIdStrictness,
  21                           setIdStrictness, setInlinePragma, idWantsToBeINLINEd,
  22                           setIdWorkerInfo, getIdCprInfo )
  23 import VarSet
  24 import Type             ( splitAlgTyConApp_maybe )
  25 import IdInfo           ( mkStrictnessInfo, noStrictnessInfo, StrictnessInfo(..),
  26                           InlinePragInfo(..), CprInfo(..) )
  27 import Demand           ( wwLazy )
  28 import SaLib
  29 import UniqSupply       ( UniqSupply, initUs, returnUs, thenUs, mapUs, getUniqueUs, UniqSM )
  30 import UniqSet
  31 import WwLib
  32 import Outputable
  33 \end{code}
  34
  35 We take Core bindings whose binders have:
  36
  37 \begin{enumerate}
  38
  39 \item Strictness attached (by the front-end of the strictness
  40 analyser), and / or
  41
  42 \item Constructed Product Result information attached by the CPR
  43 analysis pass.
  44
  45 \end{enumerate}
  46
  47 and we return some ``plain'' bindings which have been
  48 worker/wrapper-ified, meaning:
  49
  50 \begin{enumerate}
  51
  52 \item Functions have been split into workers and wrappers where
  53 appropriate.  If a function has both strictness and CPR properties
  54 then only one worker/wrapper doing both transformations is produced;
  55
  56 \item Binders' @IdInfos@ have been updated to reflect the existence of
  57 these workers/wrappers (this is where we get STRICTNESS and CPR pragma
  58 info for exported values).
  59 \end{enumerate}
  60
  61 \begin{code}
  62
  63 wwTopBinds :: UniqSupply
  64              -> [CoreBind]
  65              -> IO [CoreBind]
  66
  67 wwTopBinds us binds
  68   = do {
  69         beginPass "Worker Wrapper binds";
  70
  71         -- Create worker/wrappers, and mark binders with their
  72         -- "strictness info" [which encodes their worker/wrapper-ness]
  73         let { binds' = workersAndWrappers us binds };
  74
  75         endPass "Worker Wrapper binds" (opt_D_dump_worker_wrapper ||
  76                                         opt_D_verbose_core2core) binds'
  77     }
  78 \end{code}
  79
  80
  81 \begin{code}
  82 workersAndWrappers :: UniqSupply -> [CoreBind] -> [CoreBind]
  83
  84 workersAndWrappers us top_binds
  85   = initUs us $
  86     mapUs (wwBind True{-top-level-}) top_binds `thenUs` \ top_binds2 ->
  87     let
  88         top_binds3 = map make_top_binding top_binds2
  89     in
  90     returnUs (concat top_binds3)
  91   where
  92     make_top_binding :: WwBinding -> [CoreBind]
  93
  94     make_top_binding (WwLet binds) = binds
  95 \end{code}
  96
  97 %************************************************************************
  98 %*                                                                      *
  99 \subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
 100 %*                                                                      *
 101 %************************************************************************
 102
 103 @wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
 104 turn.  Non-recursive case first, then recursive...
 105
 106 \begin{code}
 107 wwBind  :: Bool                 -- True <=> top-level binding
 108         -> CoreBind
 109         -> UniqSM WwBinding     -- returns a WwBinding intermediate form;
 110                                 -- the caller will convert to Expr/Binding,
 111                                 -- as appropriate.
 112
 113 wwBind top_level (NonRec binder rhs)
 114   = wwExpr rhs                                          `thenUs` \ new_rhs ->
 115     tryWW True {- non-recursive -} binder new_rhs       `thenUs` \ new_pairs ->
 116     returnUs (WwLet [NonRec b e | (b,e) <- new_pairs])
 117       -- Generated bindings must be non-recursive
 118       -- because the original binding was.
 119
 120 ------------------------------
 121
 122 wwBind top_level (Rec pairs)
 123   = mapUs do_one pairs          `thenUs` \ new_pairs ->
 124     returnUs (WwLet [Rec (concat new_pairs)])
 125   where
 126     do_one (binder, rhs) = wwExpr rhs   `thenUs` \ new_rhs ->
 127                            tryWW False {- recursive -} binder new_rhs
 128 \end{code}
 129
 130 @wwExpr@ basically just walks the tree, looking for appropriate
 131 annotations that can be used. Remember it is @wwBind@ that does the
 132 matching by looking for strict arguments of the correct type.
 133 @wwExpr@ is a version that just returns the ``Plain'' Tree.
 134 ???????????????? ToDo
 135
 136 \begin{code}
 137 wwExpr :: CoreExpr -> UniqSM CoreExpr
 138
 139 wwExpr e@(Type _)   = returnUs e
 140 wwExpr e@(Var _)    = returnUs e
 141
 142 wwExpr e@(Con con args)
 143  = mapUs wwExpr args    `thenUs` \ args' ->
 144    returnUs (Con con args')
 145
 146 wwExpr (Lam binder expr)
 147   = wwExpr expr                 `thenUs` \ new_expr ->
 148     returnUs (Lam binder new_expr)
 149
 150 wwExpr (App f a)
 151   = wwExpr f                    `thenUs` \ new_f ->
 152     wwExpr a                    `thenUs` \ new_a ->
 153     returnUs (App new_f new_a)
 154
 155 wwExpr (Note note expr)
 156   = wwExpr expr                 `thenUs` \ new_expr ->
 157     returnUs (Note note new_expr)
 158
 159 wwExpr (Let bind expr)
 160   = wwBind False{-not top-level-} bind  `thenUs` \ intermediate_bind ->
 161     wwExpr expr                         `thenUs` \ new_expr ->
 162     returnUs (mash_ww_bind intermediate_bind new_expr)
 163   where
 164     mash_ww_bind (WwLet  binds)   body = mkLets binds body
 165     mash_ww_bind (WwCase case_fn) body = case_fn body
 166
 167 wwExpr (Case expr binder alts)
 168   = wwExpr expr                         `thenUs` \ new_expr ->
 169     mapUs ww_alt alts                   `thenUs` \ new_alts ->
 170     returnUs (Case new_expr binder new_alts)
 171   where
 172     ww_alt (con, binders, rhs)
 173       = wwExpr rhs                      `thenUs` \ new_rhs ->
 174         returnUs (con, binders, new_rhs)
 175 \end{code}
 176
 177 %************************************************************************
 178 %*                                                                      *
 179 \subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
 180 %*                                                                      *
 181 %************************************************************************
 182
 183 @tryWW@ just accumulates arguments, converts strictness info from the
 184 front-end into the proper form, then calls @mkWwBodies@ to do
 185 the business.
 186
 187 We have to BE CAREFUL that we don't worker-wrapperize an Id that has
 188 already been w-w'd!  (You can end up with several liked-named Ids
 189 bouncing around at the same time---absolute mischief.)  So the
 190 criterion we use is: if an Id already has an unfolding (for whatever
 191 reason), then we don't w-w it.
 192
 193 The only reason this is monadised is for the unique supply.
 194
 195 \begin{code}
 196 tryWW   :: Bool                         -- True <=> a non-recursive binding
 197         -> Id                           -- The fn binder
 198         -> CoreExpr                     -- The bound rhs; its innards
 199                                         --   are already ww'd
 200         -> UniqSM [(Id, CoreExpr)]      -- either *one* or *two* pairs;
 201                                         -- if one, then no worker (only
 202                                         -- the orig "wrapper" lives on);
 203                                         -- if two, then a worker and a
 204                                         -- wrapper.
 205 tryWW non_rec fn_id rhs
 206   |  idWantsToBeINLINEd fn_id
 207   || (non_rec &&        -- Don't split if its non-recursive and small
 208       certainlySmallEnoughToInline fn_id unfold_guidance
 209      )
 210             -- No point in worker/wrappering something that is going to be
 211             -- INLINEd wholesale anyway.  If the strictness analyser is run
 212             -- twice, this test also prevents wrappers (which are INLINEd)
 213             -- from being re-done.
 214
 215   || not (do_strict_ww || do_cpr_ww)
 216   = returnUs [ (fn_id, rhs) ]
 217
 218   | otherwise           -- Do w/w split
 219   = let
 220         (tyvars, wrap_args, body) = collectTyAndValBinders rhs
 221     in
 222     mkWwBodies tyvars wrap_args
 223                (coreExprType body)
 224                revised_wrap_args_info
 225                cpr_info
 226                                                 `thenUs` \ (wrap_fn, work_fn, work_demands) ->
 227     getUniqueUs                                 `thenUs` \ work_uniq ->
 228     let
 229         work_rhs  = work_fn body
 230         work_id   = mkWorkerId work_uniq fn_id (coreExprType work_rhs) `setIdStrictness`
 231                     (if do_strict_ww then mkStrictnessInfo (work_demands, result_bot)
 232                                      else noStrictnessInfo)
 233
 234         wrap_rhs = wrap_fn work_id
 235         wrap_id  = fn_id `setIdStrictness`
 236                          (if do_strict_ww then mkStrictnessInfo (revised_wrap_args_info, result_bot)
 237                                           else noStrictnessInfo)
 238                          `setIdWorkerInfo` (Just work_id)
 239                          `setInlinePragma` IWantToBeINLINEd
 240                 -- Add info to the wrapper:
 241                 --      (a) we want to inline it everywhere
 242                 --      (b) we want to pin on its revised strictness info
 243                 --      (c) we pin on its worker id
 244     in
 245     returnUs ([(work_id, work_rhs), (wrap_id, wrap_rhs)])
 246         -- Worker first, because wrapper mentions it
 247   where
 248     strictness_info     = getIdStrictness fn_id
 249     has_strictness_info = case strictness_info of
 250                                 StrictnessInfo _ _ -> True
 251                                 other              -> False
 252
 253     StrictnessInfo wrap_args_info result_bot = strictness_info
 254
 255     revised_wrap_args_info = if has_strictness_info
 256                                then setUnpackStrategy wrap_args_info
 257                                else repeat wwLazy
 258
 259
 260     -- If we are going to split for CPR purposes anyway,  then
 261     -- we may as well do the strictness transformation
 262     do_strict_ww = has_strictness_info && (do_cpr_ww ||
 263                                            worthSplitting revised_wrap_args_info)
 264
 265     cpr_info     = getIdCprInfo fn_id
 266     has_cpr_info = case cpr_info of
 267                                 CPRInfo _ -> True
 268                                 other     -> False
 269
 270     do_cpr_ww = has_cpr_info
 271
 272     unfold_guidance = calcUnfoldingGuidance opt_UnfoldingCreationThreshold rhs
 273
 274 -- This rather (nay! extremely!) crude function looks at a wrapper function, and
 275 -- snaffles out (a) the worker Id and (b) constructors needed to
 276 -- make the wrapper.
 277 -- These are needed when we write an interface file.
 278
 279 -- <Mar 1999 (keving)> - Well,  since the addition of the CPR transformation this function
 280 -- got too crude!
 281 -- Now the worker id is stored directly in the id's Info field.  We still use this function to
 282 -- snaffle the wrapper's constructors but I don't trust the code to find the worker id.
 283 getWorkerIdAndCons :: Id -> CoreExpr -> (Id, UniqSet DataCon)
 284 getWorkerIdAndCons wrap_id wrapper_fn
 285   = (work_id wrapper_fn, get_cons wrapper_fn)
 286   where
 287
 288     work_id wrapper_fn
 289             = case get_work_id wrapper_fn of
 290                 []   -> case work_id_try2 wrapper_fn of
 291                         [] -> pprPanic "getWorkerIdAndCons: can't find worker id" (ppr wrap_id)
 292                         [id] -> id
 293                         _    -> pprPanic "getWorkerIdAndCons: found too many worker ids" (ppr wrap_id)
 294                 [id] -> id
 295                 _    -> pprPanic "getWorkerIdAndCons: found too many worker ids" (ppr wrap_id)
 296
 297     get_work_id (Lam _ body)                     = get_work_id body
 298     get_work_id (Case _ _ [(_,_,rhs@(Case _ _ _))])     = get_work_id rhs
 299     get_work_id (Case scrut _ [(_,_,rhs)])              = (get_work_id scrut) ++ (get_work_id rhs)
 300     get_work_id (Note _ body)                    = get_work_id body
 301     get_work_id (Let _ body)                     = get_work_id body
 302     get_work_id (App (Var work_id) _)            = [work_id]
 303     get_work_id (App fn _)                       = get_work_id fn
 304     get_work_id (Var work_id)                    = []
 305     get_work_id other                            = []
 306
 307     work_id_try2 (Lam _ body)                    = work_id_try2 body
 308     work_id_try2 (Note _ body)                   = work_id_try2 body
 309     work_id_try2 (Let _ body)                    = work_id_try2 body
 310     work_id_try2 (App fn _)                      = work_id_try2 fn
 311     work_id_try2 (Var work_id)                   = [work_id]
 312     work_id_try2 other                           = []
 313
 314     get_cons (Lam _ body)                       = get_cons body
 315     get_cons (Let (NonRec _ rhs) body)          = get_cons rhs `unionUniqSets` get_cons body
 316
 317     get_cons (Case e _ [(DataCon dc,_,rhs)])    = (get_cons e `unionUniqSets` get_cons rhs)
 318                                                   `addOneToUniqSet` dc
 319
 320         -- Coercions don't mention the construtor now,
 321         -- but we must still put the constructor in the interface
 322         -- file so that the RHS of the newtype decl is imported
 323     get_cons (Note (Coerce to_ty from_ty) body)
 324         = get_cons body `addOneToUniqSet` con
 325         where
 326           con = case splitAlgTyConApp_maybe from_ty of
 327                         Just (_, _, [con]) -> con
 328                         other              -> pprPanic "getWorkerIdAndCons" (ppr to_ty)
 329
 330     get_cons other = emptyUniqSet
 331 \end{code}