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 ( 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 Const            ( Con(..) )
  17 import DataCon          ( DataCon )
  18 import MkId             ( mkWorkerId )
  19 import Id               ( Id, getIdStrictness,
  20                           setIdStrictness, setInlinePragma, idWantsToBeINLINEd,
  21                         )
  22 import VarSet
  23 import Type             ( splitAlgTyConApp_maybe )
  24 import IdInfo           ( mkStrictnessInfo, StrictnessInfo(..),
  25                           InlinePragInfo(..) )
  26 import SaLib
  27 import UniqSupply       ( UniqSupply, initUs, returnUs, thenUs, mapUs, getUniqueUs, UniqSM )
  28 import UniqSet
  29 import WwLib
  30 import Outputable
  31 \end{code}
  32
  33 We take Core bindings whose binders have their strictness attached (by
  34 the front-end of the strictness analyser), and we return some
  35 ``plain'' bindings which have been worker/wrapper-ified, meaning:
  36 \begin{enumerate}
  37 \item
  38 Functions have been split into workers and wrappers where appropriate;
  39 \item
  40 Binders' @IdInfos@ have been updated to reflect the existence
  41 of these workers/wrappers (this is where we get STRICTNESS pragma
  42 info for exported values).
  43 \end{enumerate}
  44
  45 \begin{code}
  46 workersAndWrappers :: UniqSupply -> [CoreBind] -> [CoreBind]
  47
  48 workersAndWrappers us top_binds
  49   = initUs us $
  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 -> [CoreBind]
  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         -> CoreBind
  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 True {- non-recursive -} 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 False {- recursive -} 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@(Type _)   = returnUs e
 104 wwExpr e@(Var _)    = returnUs e
 105
 106 wwExpr e@(Con con args)
 107  = mapUs wwExpr args    `thenUs` \ args' ->
 108    returnUs (Con con args')
 109
 110 wwExpr (Lam binder expr)
 111   = wwExpr expr                 `thenUs` \ new_expr ->
 112     returnUs (Lam binder new_expr)
 113
 114 wwExpr (App f a)
 115   = wwExpr f                    `thenUs` \ new_f ->
 116     wwExpr a                    `thenUs` \ new_a ->
 117     returnUs (App new_f new_a)
 118
 119 wwExpr (Note note expr)
 120   = wwExpr expr                 `thenUs` \ new_expr ->
 121     returnUs (Note note new_expr)
 122
 123 wwExpr (Let bind expr)
 124   = wwBind False{-not top-level-} bind  `thenUs` \ intermediate_bind ->
 125     wwExpr expr                         `thenUs` \ new_expr ->
 126     returnUs (mash_ww_bind intermediate_bind new_expr)
 127   where
 128     mash_ww_bind (WwLet  binds)   body = mkLets binds body
 129     mash_ww_bind (WwCase case_fn) body = case_fn body
 130
 131 wwExpr (Case expr binder alts)
 132   = wwExpr expr                         `thenUs` \ new_expr ->
 133     mapUs ww_alt alts                   `thenUs` \ new_alts ->
 134     returnUs (Case new_expr binder new_alts)
 135   where
 136     ww_alt (con, binders, rhs)
 137       = wwExpr rhs                      `thenUs` \ new_rhs ->
 138         returnUs (con, binders, new_rhs)
 139 \end{code}
 140
 141 %************************************************************************
 142 %*                                                                      *
 143 \subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
 144 %*                                                                      *
 145 %************************************************************************
 146
 147 @tryWW@ just accumulates arguments, converts strictness info from the
 148 front-end into the proper form, then calls @mkWwBodies@ to do
 149 the business.
 150
 151 We have to BE CAREFUL that we don't worker-wrapperize an Id that has
 152 already been w-w'd!  (You can end up with several liked-named Ids
 153 bouncing around at the same time---absolute mischief.)  So the
 154 criterion we use is: if an Id already has an unfolding (for whatever
 155 reason), then we don't w-w it.
 156
 157 The only reason this is monadised is for the unique supply.
 158
 159 \begin{code}
 160 tryWW   :: Bool                         -- True <=> a non-recursive binding
 161         -> Id                           -- The fn binder
 162         -> CoreExpr                     -- The bound rhs; its innards
 163                                         --   are already ww'd
 164         -> UniqSM [(Id, CoreExpr)]      -- either *one* or *two* pairs;
 165                                         -- if one, then no worker (only
 166                                         -- the orig "wrapper" lives on);
 167                                         -- if two, then a worker and a
 168                                         -- wrapper.
 169 tryWW non_rec fn_id rhs
 170   |  idWantsToBeINLINEd fn_id
 171   || (non_rec &&        -- Don't split if its non-recursive and small
 172       certainlySmallEnoughToInline fn_id unfold_guidance
 173      )
 174             -- No point in worker/wrappering something that is going to be
 175             -- INLINEd wholesale anyway.  If the strictness analyser is run
 176             -- twice, this test also prevents wrappers (which are INLINEd)
 177             -- from being re-done.
 178
 179   || not has_strictness_info
 180   || not (worthSplitting revised_wrap_args_info)
 181   = returnUs [ (fn_id, rhs) ]
 182
 183   | otherwise           -- Do w/w split
 184   = let
 185         (tyvars, wrap_args, body) = collectTyAndValBinders rhs
 186     in
 187     mkWwBodies tyvars wrap_args
 188                (coreExprType body)
 189                revised_wrap_args_info           `thenUs` \ (wrap_fn, work_fn, work_demands) ->
 190     getUniqueUs                                 `thenUs` \ work_uniq ->
 191     let
 192         work_rhs  = work_fn body
 193         work_id   = mkWorkerId work_uniq fn_id (coreExprType work_rhs) `setIdStrictness`
 194                     mkStrictnessInfo work_demands False
 195
 196         wrap_rhs = wrap_fn work_id
 197         wrap_id  = fn_id `setIdStrictness` mkStrictnessInfo revised_wrap_args_info True
 198                          `setInlinePragma` IWantToBeINLINEd
 199                 -- Add info to the wrapper:
 200                 --      (a) we want to inline it everywhere
 201                 --      (b) we want to pin on its revised stricteness info
 202                 --      (c) we pin on its worker id and the list of constructors mentioned in the wrapper
 203     in
 204     returnUs ([(work_id, work_rhs), (wrap_id, wrap_rhs)])
 205         -- Worker first, because wrapper mentions it
 206   where
 207     strictness_info     = getIdStrictness fn_id
 208     has_strictness_info = case strictness_info of
 209                                 StrictnessInfo _ _ -> True
 210                                 other              -> False
 211
 212     wrap_args_info = case strictness_info of
 213                         StrictnessInfo args_info _ -> args_info
 214     revised_wrap_args_info = setUnpackStrategy wrap_args_info
 215
 216     unfold_guidance = calcUnfoldingGuidance opt_UnfoldingCreationThreshold rhs
 217
 218 -- This rather (nay! extremely!) crude function looks at a wrapper function, and
 219 -- snaffles out (a) the worker Id and (b) constructors needed to
 220 -- make the wrapper.
 221 -- These are needed when we write an interface file.
 222 getWorkerIdAndCons :: Id -> CoreExpr -> (Id, UniqSet DataCon)
 223 getWorkerIdAndCons wrap_id wrapper_fn
 224   = (get_work_id wrapper_fn, get_cons wrapper_fn)
 225   where
 226     get_work_id (Lam _ body)                     = get_work_id body
 227     get_work_id (Case _ _ [(_,_,rhs)])           = get_work_id rhs
 228     get_work_id (Note _ body)                    = get_work_id body
 229     get_work_id (Let _ body)                     = get_work_id body
 230     get_work_id (App fn _)                       = get_work_id fn
 231     get_work_id (Var work_id)                    = work_id
 232     get_work_id other                            = pprPanic "getWorkerIdAndCons" (ppr wrap_id)
 233
 234
 235     get_cons (Lam _ body)                       = get_cons body
 236     get_cons (Let (NonRec _ rhs) body)          = get_cons rhs `unionUniqSets` get_cons body
 237
 238     get_cons (Case e _ [(DataCon dc,_,rhs)])    = (get_cons e `unionUniqSets` get_cons rhs)
 239                                                   `addOneToUniqSet` dc
 240
 241         -- Coercions don't mention the construtor now,
 242         -- but we must still put the constructor in the interface
 243         -- file so that the RHS of the newtype decl is imported
 244     get_cons (Note (Coerce to_ty from_ty) body)
 245         = get_cons body `addOneToUniqSet` con
 246         where
 247           con = case splitAlgTyConApp_maybe from_ty of
 248                         Just (_, _, [con]) -> con
 249                         other              -> pprPanic "getWorkerIdAndCons" (ppr to_ty)
 250
 251     get_cons other = emptyUniqSet
 252 \end{code}