ghc/compiler/stranal/WwLib.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1996
   3 %
   4 \section[WwLib]{A library for the ``worker/wrapper'' back-end to the strictness analyser}
   5
   6 \begin{code}
   7 #include "HsVersions.h"
   8
   9 module WwLib (
  10         WwBinding(..),
  11
  12         worthSplitting, setUnpackStrategy,
  13         mkWwBodies, mkWrapper
  14     ) where
  15
  16 IMP_Ubiq(){-uitous-}
  17
  18 import CoreSyn
  19 import Id               ( idType, mkSysLocal, dataConArgTys )
  20 import IdInfo           ( mkStrictnessInfo, nonAbsentArgs, Demand(..) )
  21 import PrelVals         ( aBSENT_ERROR_ID, voidId )
  22 import TysPrim          ( voidTy )
  23 import SrcLoc           ( noSrcLoc )
  24 import Type             ( isPrimType, mkTyVarTys, mkForAllTys, mkFunTys,
  25                           splitForAllTy, splitFunTyExpandingDicts,
  26                           maybeAppDataTyConExpandingDicts
  27                         )
  28 import UniqSupply       ( returnUs, thenUs, thenMaybeUs,
  29                           getUniques, getUnique, SYN_IE(UniqSM)
  30                         )
  31 import Util             ( zipWithEqual, zipEqual, assertPanic, panic )
  32 \end{code}
  33
  34 %************************************************************************
  35 %*                                                                      *
  36 \subsection[datatype-WwLib]{@WwBinding@: a datatype for worker/wrapper-ing}
  37 %*                                                                      *
  38 %************************************************************************
  39
  40 In the worker/wrapper stuff, we want to carry around @CoreBindings@ in
  41 an ``intermediate form'' that can later be turned into a \tr{let} or
  42 \tr{case} (depending on strictness info).
  43
  44 \begin{code}
  45 data WwBinding
  46   = WwLet  [CoreBinding]
  47   | WwCase (CoreExpr -> CoreExpr)
  48                 -- the "case" will be a "strict let" of the form:
  49                 --
  50                 --  case rhs of
  51                 --    <blah> -> body
  52                 --
  53                 -- (instead of "let <blah> = rhs in body")
  54                 --
  55                 -- The expr you pass to the function is "body" (the
  56                 -- expression that goes "in the corner").
  57 \end{code}
  58
  59 %************************************************************************
  60 %*                                                                      *
  61 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
  62 %*                                                                      *
  63 %************************************************************************
  64
  65         ************   WARNING  ******************
  66         these comments are rather out of date
  67         *****************************************
  68
  69 @mkWrapperAndWorker@ is given:
  70 \begin{enumerate}
  71 \item
  72 The {\em original function} \tr{f}, of the form:
  73 \begin{verbatim}
  74 f = /\ tyvars -> \ args -> body
  75 \end{verbatim}
  76 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
  77 are given separately.
  78
  79 We use the Id \tr{f} mostly to get its type.
  80
  81 \item
  82 Strictness information about \tr{f}, in the form of a list of
  83 @Demands@.
  84
  85 \item
  86 A @UniqueSupply@.
  87 \end{enumerate}
  88
  89 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
  90 \begin{enumerate}
  91 \item
  92 Maybe @Nothing@: no worker/wrappering going on in this case. This can
  93 happen (a)~if the strictness info says that there is nothing
  94 interesting to do or (b)~if *any* of the argument types corresponding
  95 to ``active'' arg postitions is abstract or will be to the outside
  96 world (i.e., {\em this} module can see the constructors, but nobody
  97 else will be able to).  An ``active'' arg position is one which the
  98 wrapper has to unpack.  An importing module can't do this unpacking,
  99 so it simply has to give up and call the wrapper only.
 100
 101 \item
 102 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
 103
 104 The @wrapper_Id@ is just the one that was passed in, with its
 105 strictness IdInfo updated.
 106 \end{enumerate}
 107
 108 The \tr{body} of the original function may not be given (i.e., it's
 109 BOTTOM), in which case you'd jolly well better not tug on the
 110 worker-body output!
 111
 112 Here's an example.  The original function is:
 113 \begin{verbatim}
 114 g :: forall a . Int -> [a] -> a
 115
 116 g = /\ a -> \ x ys ->
 117         case x of
 118           0 -> head ys
 119           _ -> head (tail ys)
 120 \end{verbatim}
 121
 122 From this, we want to produce:
 123 \begin{verbatim}
 124 -- wrapper (an unfolding)
 125 g :: forall a . Int -> [a] -> a
 126
 127 g = /\ a -> \ x ys ->
 128         case x of
 129           I# x# -> g.wrk a x# ys
 130             -- call the worker; don't forget the type args!
 131
 132 -- worker
 133 g.wrk :: forall a . Int# -> [a] -> a
 134
 135 g.wrk = /\ a -> \ x# ys ->
 136         let
 137             x = I# x#
 138         in
 139             case x of               -- note: body of g moved intact
 140               0 -> head ys
 141               _ -> head (tail ys)
 142 \end{verbatim}
 143
 144 Something we have to be careful about:  Here's an example:
 145 \begin{verbatim}
 146 -- "f" strictness: U(P)U(P)
 147 f (I# a) (I# b) = a +# b
 148
 149 g = f   -- "g" strictness same as "f"
 150 \end{verbatim}
 151 \tr{f} will get a worker all nice and friendly-like; that's good.
 152 {\em But we don't want a worker for \tr{g}}, even though it has the
 153 same strictness as \tr{f}.  Doing so could break laziness, at best.
 154
 155 Consequently, we insist that the number of strictness-info items is
 156 exactly the same as the number of lambda-bound arguments.  (This is
 157 probably slightly paranoid, but OK in practice.)  If it isn't the
 158 same, we ``revise'' the strictness info, so that we won't propagate
 159 the unusable strictness-info into the interfaces.
 160
 161
 162 %************************************************************************
 163 %*                                                                      *
 164 \subsection{Functions over Demands}
 165 %*                                                                      *
 166 %************************************************************************
 167
 168 \begin{code}
 169 mAX_WORKER_ARGS :: Int          -- ToDo: set via flag
 170 mAX_WORKER_ARGS = 6
 171
 172 setUnpackStrategy :: [Demand] -> [Demand]
 173 setUnpackStrategy ds
 174   = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
 175   where
 176     go :: Int                   -- Max number of args available for sub-components of [Demand]
 177        -> [Demand]
 178        -> (Int, [Demand])       -- Args remaining after subcomponents of [Demand] are unpacked
 179
 180     go n (WwUnpack _ cs : ds) | n' >= 0
 181                               = WwUnpack True cs' `cons` go n'' ds
 182                               | otherwise
 183                               = WwUnpack False cs `cons` go n ds
 184                               where
 185                                 n' = n + 1 - nonAbsentArgs cs
 186                                         -- Add one because we don't pass the top-level arg any more
 187                                         -- Delete # of non-absent args to which we'll now be committed
 188                                 (n'',cs') = go n' cs
 189
 190     go n (d:ds) = d `cons` go n ds
 191     go n []     = (n,[])
 192
 193     cons d (n,ds) = (n, d:ds)
 194
 195 nonAbsentArgs :: [Demand] -> Int
 196 nonAbsentArgs []                 = 0
 197 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
 198 nonAbsentArgs (d           : ds) = 1 + nonAbsentArgs ds
 199
 200 worthSplitting :: [Demand] -> Bool      -- True <=> the wrapper would not be an identity function
 201 worthSplitting []                       = False
 202 worthSplitting (WwLazy True : ds)       = True          -- Absent arg
 203 worthSplitting (WwUnpack True _ : ds)   = True          -- Arg to unpack
 204 worthSplitting (d : ds)                 = worthSplitting ds
 205
 206 allAbsent :: [Demand] -> Bool
 207 allAbsent (WwLazy True      : ds) = allAbsent ds
 208 allAbsent (WwUnpack True cs : ds) = allAbsent cs && allAbsent ds
 209 allAbsent (d                : ds) = False
 210 allAbsent []                      = True
 211 \end{code}
 212
 213
 214 %************************************************************************
 215 %*                                                                      *
 216 \subsection{The worker wrapper core}
 217 %*                                                                      *
 218 %************************************************************************
 219
 220 @mkWrapper@ is called when importing a function.  We have the type of
 221 the function and the name of its worker, and we want to make its body (the wrapper).
 222
 223 \begin{code}
 224 mkWrapper :: Type               -- Wrapper type
 225           -> [Demand]           -- Wrapper strictness info
 226           -> UniqSM (Id -> CoreExpr)    -- Wrapper body, missing worker Id
 227
 228 mkWrapper fun_ty demands
 229   = let
 230         n_wrap_args = length demands
 231     in
 232     getUniques n_wrap_args      `thenUs` \ wrap_uniqs ->
 233     let
 234         (tyvars, tau_ty)   = splitForAllTy fun_ty
 235         (arg_tys, body_ty) = splitFunTyExpandingDicts tau_ty
 236         wrap_args          = zipWith mk_ww_local wrap_uniqs arg_tys
 237         leftover_arg_tys   = drop n_wrap_args arg_tys
 238         final_body_ty      = mkFunTys leftover_arg_tys body_ty
 239     in
 240     mkWwBodies tyvars wrap_args final_body_ty demands   `thenUs` \ (wrap_fn, _, _) ->
 241     returnUs wrap_fn
 242 \end{code}
 243
 244 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
 245
 246 \begin{code}
 247 mkWwBodies :: [TyVar] -> [Id] -> Type           -- Original fn args and body type
 248            -> [Demand]                          -- Strictness info for original fn; corresp 1-1 with args
 249            -> UniqSM (Id -> CoreExpr,           -- Wrapper body, lacking only the worker Id
 250                       CoreExpr -> CoreExpr,     -- Worker body, lacking the original function body
 251                       [Demand])                 -- Strictness info for worker
 252
 253 mkWwBodies tyvars args body_ty demands
 254   | allAbsent demands &&
 255     isPrimType body_ty
 256   =     -- Horrid special case.  If the worker would have no arguments, and the
 257         -- function returns a primitive type value, that would make the worker into
 258         -- an unboxed value.  We box it by passing a dummy void argument, thus:
 259         --
 260         --      f = /\abc. \xyz. fw abc void
 261         --      fw = /\abc. \v. body
 262         --
 263     getUnique           `thenUs` \ void_arg_uniq ->
 264     let
 265         void_arg = mk_ww_local void_arg_uniq voidTy
 266     in
 267     returnUs (\ work_id -> mkLam tyvars args (App (mkTyApp (Var work_id) (mkTyVarTys tyvars)) (VarArg voidId)),
 268               \ body    -> mkLam tyvars [void_arg] body,
 269               [WwLazy True])
 270
 271 mkWwBodies tyvars args body_ty demands
 272   | otherwise
 273   = let
 274         args_w_demands = zipEqual "mkWwBodies" args demands
 275     in
 276     mkWW args_w_demands         `thenUs` \ (wrap_fn, work_args_w_demands, work_fn) ->
 277     let
 278         (work_args, work_demands) = unzip work_args_w_demands
 279     in
 280     returnUs (\ work_id -> mkLam tyvars args (wrap_fn (mkTyApp (Var work_id) (mkTyVarTys tyvars))),
 281               \ body    -> mkLam tyvars work_args (work_fn body),
 282               work_demands)
 283 \end{code}
 284
 285
 286 \begin{code}
 287 mkWW :: [(Id,Demand)]
 288      -> UniqSM (CoreExpr -> CoreExpr,   -- Wrapper body, lacking the inner call to the worker
 289                                         -- and without its lambdas
 290                 [(Id,Demand)],          -- Worker args and their demand infos
 291                 CoreExpr -> CoreExpr)   -- Worker body, lacking the original body of the function
 292
 293
 294         -- Empty case
 295 mkWW []
 296   = returnUs (\ wrapper_body -> wrapper_body,
 297               [],
 298               \ worker_body  -> worker_body)
 299
 300
 301         -- Absent case
 302 mkWW ((arg,WwLazy True) : ds)
 303   = mkWW ds             `thenUs` \ (wrap_fn, worker_args, work_fn) ->
 304     returnUs (\ wrapper_body -> wrap_fn wrapper_body,
 305               worker_args,
 306               \ worker_body  -> mk_absent_let arg (work_fn worker_body))
 307
 308
 309         -- Unpack case
 310 mkWW ((arg,WwUnpack True cs) : ds)
 311   = getUniques (length inst_con_arg_tys)                `thenUs` \ uniqs ->
 312     let
 313         unpk_args        = zipWith mk_ww_local uniqs inst_con_arg_tys
 314         unpk_args_w_ds   = zipEqual "mkWW" unpk_args cs
 315     in
 316     mkWW (unpk_args_w_ds ++ ds)         `thenUs` \ (wrap_fn, worker_args, work_fn) ->
 317     returnUs (\ wrapper_body -> mk_unpk_case arg unpk_args data_con arg_tycon (wrap_fn wrapper_body),
 318               worker_args,
 319               \ worker_body  -> work_fn (mk_pk_let arg data_con tycon_arg_tys unpk_args worker_body))
 320   where
 321     inst_con_arg_tys = dataConArgTys data_con tycon_arg_tys
 322     (arg_tycon, tycon_arg_tys, data_con)
 323         = case (maybeAppDataTyConExpandingDicts (idType arg)) of
 324
 325               Just (arg_tycon, tycon_arg_tys, [data_con]) ->
 326                                      -- The main event: a single-constructor data type
 327                                      (arg_tycon, tycon_arg_tys, data_con)
 328
 329               Just (_, _, data_cons) ->  panic "mk_ww_arg_processing: not one constr"
 330               Nothing                ->  panic "mk_ww_arg_processing: not datatype"
 331
 332
 333         -- Other cases
 334 mkWW ((arg,other_demand) : ds)
 335   = mkWW ds             `thenUs` \ (wrap_fn, worker_args, work_fn) ->
 336     returnUs (\ wrapper_body -> wrap_fn (App wrapper_body (VarArg arg)),
 337               (arg,other_demand) : worker_args,
 338               work_fn)
 339 \end{code}
 340
 341
 342 %************************************************************************
 343 %*                                                                      *
 344 \subsection{Utilities}
 345 %*                                                                      *
 346 %************************************************************************
 347
 348
 349 \begin{code}
 350 mk_absent_let arg body
 351   | not (isPrimType arg_ty)
 352   = Let (NonRec arg (mkTyApp (Var aBSENT_ERROR_ID) [arg_ty])) body
 353   | otherwise
 354   = panic "WwLib: haven't done mk_absent_let for primitives yet"
 355   where
 356     arg_ty = idType arg
 357
 358 mk_unpk_case arg unpk_args boxing_con boxing_tycon body
 359   = Case (Var arg)
 360          (AlgAlts [(boxing_con, unpk_args, body)]
 361                   NoDefault
 362          )
 363
 364 mk_pk_let arg boxing_con con_tys unpk_args body
 365   = Let (NonRec arg (Con boxing_con con_args)) body
 366   where
 367     con_args = map TyArg con_tys ++ map VarArg unpk_args
 368
 369 mk_ww_local uniq ty
 370   = mkSysLocal SLIT("ww") uniq ty noSrcLoc
 371 \end{code}