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