ghc/compiler/stranal/WorkWrap.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1996
   3 %
   4 \section[WorkWrap]{Worker/wrapper-generating back-end of strictness analyser}
   5
   6 \begin{code}
   7 module WorkWrap ( workersAndWrappers, getWorkerIdAndCons ) where
   8
   9 #include "HsVersions.h"
  10
  11 import CoreSyn
  12 import CoreUnfold       ( Unfolding, certainlySmallEnoughToInline, calcUnfoldingGuidance )
  13 import CmdLineOpts      ( opt_UnfoldingCreationThreshold )
  14
  15 import CoreUtils        ( coreExprType )
  16 import MkId             ( mkWorkerId )
  17 import Id               ( getInlinePragma, getIdStrictness,
  18                           addIdStrictness, addInlinePragma, idWantsToBeINLINEd,
  19                           IdSet, emptyIdSet, addOneToIdSet,
  20                           GenId, Id
  21                         )
  22 import IdInfo           ( noIdInfo, mkStrictnessInfo, setStrictnessInfo, StrictnessInfo(..) )
  23 import SaLib
  24 import UniqSupply       ( returnUs, thenUs, mapUs, getUnique, UniqSM )
  25 import WwLib
  26 import Outputable
  27 \end{code}
  28
  29 We take Core bindings whose binders have their strictness attached (by
  30 the front-end of the strictness analyser), and we return some
  31 ``plain'' bindings which have been worker/wrapper-ified, meaning:
  32 \begin{enumerate}
  33 \item
  34 Functions have been split into workers and wrappers where appropriate;
  35 \item
  36 Binders' @IdInfos@ have been updated to reflect the existence
  37 of these workers/wrappers (this is where we get STRICTNESS pragma
  38 info for exported values).
  39 \end{enumerate}
  40
  41 \begin{code}
  42 workersAndWrappers :: [CoreBinding] -> UniqSM [CoreBinding]
  43
  44 workersAndWrappers top_binds
  45   = mapUs (wwBind True{-top-level-}) top_binds `thenUs` \ top_binds2 ->
  46     let
  47         top_binds3 = map make_top_binding top_binds2
  48     in
  49     returnUs (concat top_binds3)
  50   where
  51     make_top_binding :: WwBinding -> [CoreBinding]
  52
  53     make_top_binding (WwLet binds) = binds
  54 \end{code}
  55
  56 %************************************************************************
  57 %*                                                                      *
  58 \subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
  59 %*                                                                      *
  60 %************************************************************************
  61
  62 @wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
  63 turn.  Non-recursive case first, then recursive...
  64
  65 \begin{code}
  66 wwBind  :: Bool                 -- True <=> top-level binding
  67         -> CoreBinding
  68         -> UniqSM WwBinding     -- returns a WwBinding intermediate form;
  69                                 -- the caller will convert to Expr/Binding,
  70                                 -- as appropriate.
  71
  72 wwBind top_level (NonRec binder rhs)
  73   = wwExpr rhs                  `thenUs` \ new_rhs ->
  74     tryWW binder new_rhs        `thenUs` \ new_pairs ->
  75     returnUs (WwLet [NonRec b e | (b,e) <- new_pairs])
  76       -- Generated bindings must be non-recursive
  77       -- because the original binding was.
  78
  79 ------------------------------
  80
  81 wwBind top_level (Rec pairs)
  82   = mapUs do_one pairs          `thenUs` \ new_pairs ->
  83     returnUs (WwLet [Rec (concat new_pairs)])
  84   where
  85     do_one (binder, rhs) = wwExpr rhs   `thenUs` \ new_rhs ->
  86                            tryWW binder new_rhs
  87 \end{code}
  88
  89 @wwExpr@ basically just walks the tree, looking for appropriate
  90 annotations that can be used. Remember it is @wwBind@ that does the
  91 matching by looking for strict arguments of the correct type.
  92 @wwExpr@ is a version that just returns the ``Plain'' Tree.
  93 ???????????????? ToDo
  94
  95 \begin{code}
  96 wwExpr :: CoreExpr -> UniqSM CoreExpr
  97
  98 wwExpr e@(Var _)    = returnUs e
  99 wwExpr e@(Lit _)    = returnUs e
 100 wwExpr e@(Con  _ _) = returnUs e
 101 wwExpr e@(Prim _ _) = returnUs e
 102
 103 wwExpr (Lam binder expr)
 104   = wwExpr expr                 `thenUs` \ new_expr ->
 105     returnUs (Lam binder new_expr)
 106
 107 wwExpr (App f a)
 108   = wwExpr f                    `thenUs` \ new_f ->
 109     returnUs (App new_f a)
 110
 111 wwExpr (Note note expr)
 112   = wwExpr expr                 `thenUs` \ new_expr ->
 113     returnUs (Note note new_expr)
 114
 115 wwExpr (Let bind expr)
 116   = wwBind False{-not top-level-} bind  `thenUs` \ intermediate_bind ->
 117     wwExpr expr                         `thenUs` \ new_expr ->
 118     returnUs (mash_ww_bind intermediate_bind new_expr)
 119   where
 120     mash_ww_bind (WwLet  binds)   body = mkCoLetsNoUnboxed binds body
 121     mash_ww_bind (WwCase case_fn) body = case_fn body
 122
 123 wwExpr (Case expr alts)
 124   = wwExpr expr                         `thenUs` \ new_expr ->
 125     ww_alts alts                        `thenUs` \ new_alts ->
 126     returnUs (Case new_expr new_alts)
 127   where
 128     ww_alts (AlgAlts alts deflt)
 129       = mapUs ww_alg_alt alts           `thenUs` \ new_alts ->
 130         ww_deflt deflt                  `thenUs` \ new_deflt ->
 131         returnUs (AlgAlts new_alts new_deflt)
 132
 133     ww_alts (PrimAlts alts deflt)
 134       = mapUs ww_prim_alt alts          `thenUs` \ new_alts ->
 135         ww_deflt deflt                  `thenUs` \ new_deflt ->
 136         returnUs (PrimAlts new_alts new_deflt)
 137
 138     ww_alg_alt (con, binders, rhs)
 139       = wwExpr rhs                      `thenUs` \ new_rhs ->
 140         returnUs (con, binders, new_rhs)
 141
 142     ww_prim_alt (lit, rhs)
 143       = wwExpr rhs                      `thenUs` \ new_rhs ->
 144         returnUs (lit, new_rhs)
 145
 146     ww_deflt NoDefault
 147       = returnUs NoDefault
 148
 149     ww_deflt (BindDefault binder rhs)
 150       = wwExpr rhs                      `thenUs` \ new_rhs ->
 151         returnUs (BindDefault binder new_rhs)
 152 \end{code}
 153
 154 %************************************************************************
 155 %*                                                                      *
 156 \subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
 157 %*                                                                      *
 158 %************************************************************************
 159
 160 @tryWW@ just accumulates arguments, converts strictness info from the
 161 front-end into the proper form, then calls @mkWwBodies@ to do
 162 the business.
 163
 164 We have to BE CAREFUL that we don't worker-wrapperize an Id that has
 165 already been w-w'd!  (You can end up with several liked-named Ids
 166 bouncing around at the same time---absolute mischief.)  So the
 167 criterion we use is: if an Id already has an unfolding (for whatever
 168 reason), then we don't w-w it.
 169
 170 The only reason this is monadised is for the unique supply.
 171
 172 \begin{code}
 173 tryWW   :: Id                           -- The fn binder
 174         -> CoreExpr                     -- The bound rhs; its innards
 175                                         --   are already ww'd
 176         -> UniqSM [(Id, CoreExpr)]      -- either *one* or *two* pairs;
 177                                         -- if one, then no worker (only
 178                                         -- the orig "wrapper" lives on);
 179                                         -- if two, then a worker and a
 180                                         -- wrapper.
 181 tryWW fn_id rhs
 182   |  idWantsToBeINLINEd fn_id
 183   || (certainlySmallEnoughToInline fn_id $
 184       calcUnfoldingGuidance opt_UnfoldingCreationThreshold rhs
 185      )
 186             -- No point in worker/wrappering something that is going to be
 187             -- INLINEd wholesale anyway.  If the strictness analyser is run
 188             -- twice, this test also prevents wrappers (which are INLINEd)
 189             -- from being re-done.
 190
 191   || not has_strictness_info
 192   || not (worthSplitting revised_wrap_args_info)
 193   = returnUs [ (fn_id, rhs) ]
 194
 195   | otherwise           -- Do w/w split
 196   = let
 197         (tyvars, wrap_args, body) = collectBinders rhs
 198     in
 199     mkWwBodies tyvars wrap_args
 200                (coreExprType body)
 201                revised_wrap_args_info           `thenUs` \ (wrap_fn, work_fn, work_demands) ->
 202     getUnique                                   `thenUs` \ work_uniq ->
 203     let
 204         work_rhs  = work_fn body
 205         work_id   = mkWorkerId work_uniq fn_id (coreExprType work_rhs) work_info
 206         work_info = mkStrictnessInfo work_demands False `setStrictnessInfo` noIdInfo
 207
 208         wrap_rhs = wrap_fn work_id
 209         wrap_id  = addInlinePragma (fn_id `addIdStrictness`
 210                                     mkStrictnessInfo revised_wrap_args_info True)
 211                 -- Add info to the wrapper:
 212                 --      (a) we want to inline it everywhere
 213                 --      (b) we want to pin on its revised stricteness info
 214                 --      (c) we pin on its worker id and the list of constructors mentioned in the wrapper
 215     in
 216     returnUs ([(work_id, work_rhs), (wrap_id, wrap_rhs)])
 217         -- Worker first, because wrapper mentions it
 218   where
 219     strictness_info     = getIdStrictness fn_id
 220     has_strictness_info = case strictness_info of
 221                                 StrictnessInfo _ _ -> True
 222                                 other              -> False
 223
 224     wrap_args_info = case strictness_info of
 225                         StrictnessInfo args_info _ -> args_info
 226     revised_wrap_args_info = setUnpackStrategy wrap_args_info
 227
 228 -- This rather (nay! extremely!) crude function looks at a wrapper function, and
 229 -- snaffles out (a) the worker Id and (b) constructors needed to
 230 -- make the wrapper.
 231 -- These are needed when we write an interface file.
 232 getWorkerIdAndCons wrap_id wrapper_fn
 233   = go wrapper_fn
 234   where
 235     go (Lam _ body)                       = go body
 236     go (Case _ (AlgAlts [(con,_,rhs)] _)) = let (wrap_id, cons) = go rhs
 237                                             in  (wrap_id, cons `addOneToIdSet` con)
 238 {-
 239         -- Coercions don't mention the construtor now,
 240         -- so I don't think we need this
 241     go (Let (NonRec _ (Coerce (CoerceOut con) _ _)) body)
 242                                           = let (wrap_id, cons) = go body
 243                                             in  (wrap_id, cons `addOneToIdSet` con)
 244 -}
 245     go other                              = (get_work_id other, emptyIdSet)
 246
 247     get_work_id (App fn _)    = get_work_id fn
 248     get_work_id (Var work_id) = work_id
 249     get_work_id other         = pprPanic "getWorkerIdAndCons" (ppr wrap_id)
 250 \end{code}