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