ghc/compiler/stranal/WwLib.lhs

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