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