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