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