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         mkWwBodies, mAX_WORKER_ARGS,
  13
  14         -- our friendly worker/wrapper monad:
  15         WwM(..),
  16         returnWw, thenWw, mapWw,
  17         getUniqueWw, uniqSMtoWwM
  18
  19         -- and to make the interface self-sufficient...
  20     ) where
  21
  22 import Ubiq{-uitous-}
  23
  24 import PrelInfo         ( aBSENT_ERROR_ID )
  25 {-
  26 import Id               ( mkWorkerId, mkSysLocal, idType,
  27                           getInstantiatedDataConSig, getIdInfo,
  28                           replaceIdInfo, addIdStrictness, DataCon(..)
  29                         )
  30 import IdInfo           -- lots of things
  31 import Maybes           ( maybeToBool, Maybe(..), MaybeErr )
  32 import SaLib
  33 import SrcLoc           ( mkUnknownSrcLoc )
  34 import Type             ( mkTyVarTy, mkFunTys, isPrimType,
  35                           maybeDataTyCon, quantifyTy
  36                         )
  37 import UniqSupply
  38 -}
  39 import Util             ( panic )
  40
  41 infixr 9 `thenWw`
  42
  43 quantifyTy = panic "WwLib.quantifyTy"
  44 \end{code}
  45
  46 %************************************************************************
  47 %*                                                                      *
  48 \subsection[datatype-WwLib]{@WwBinding@: a datatype for worker/wrapper-ing}
  49 %*                                                                      *
  50 %************************************************************************
  51
  52 In the worker/wrapper stuff, we want to carry around @CoreBindings@ in
  53 an ``intermediate form'' that can later be turned into a \tr{let} or
  54 \tr{case} (depending on strictness info).
  55
  56 \begin{code}
  57 data WwBinding
  58   = WwLet  [CoreBinding]
  59   | WwCase (CoreExpr -> CoreExpr)
  60                 -- the "case" will be a "strict let" of the form:
  61                 --
  62                 --  case rhs of
  63                 --    <blah> -> body
  64                 --
  65                 -- (instead of "let <blah> = rhs in body")
  66                 --
  67                 -- The expr you pass to the function is "body" (the
  68                 -- expression that goes "in the corner").
  69 \end{code}
  70
  71 %************************************************************************
  72 %*                                                                      *
  73 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
  74 %*                                                                      *
  75 %************************************************************************
  76
  77         ************   WARNING  ******************
  78         these comments are rather out of date
  79         *****************************************
  80
  81 @mkWrapperAndWorker@ is given:
  82 \begin{enumerate}
  83 \item
  84 The {\em original function} \tr{f}, of the form:
  85 \begin{verbatim}
  86 f = /\ tyvars -> \ args -> body
  87 \end{verbatim}
  88 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
  89 are given separately.
  90
  91 We use the Id \tr{f} mostly to get its type.
  92
  93 \item
  94 Strictness information about \tr{f}, in the form of a list of
  95 @Demands@.
  96
  97 \item
  98 A @UniqueSupply@.
  99 \end{enumerate}
 100
 101 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
 102 \begin{enumerate}
 103 \item
 104 Maybe @Nothing@: no worker/wrappering going on in this case. This can
 105 happen (a)~if the strictness info says that there is nothing
 106 interesting to do or (b)~if *any* of the argument types corresponding
 107 to ``active'' arg postitions is abstract or will be to the outside
 108 world (i.e., {\em this} module can see the constructors, but nobody
 109 else will be able to).  An ``active'' arg position is one which the
 110 wrapper has to unpack.  An importing module can't do this unpacking,
 111 so it simply has to give up and call the wrapper only.
 112
 113 \item
 114 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
 115
 116 The @wrapper_Id@ is just the one that was passed in, with its
 117 strictness IdInfo updated.
 118 \end{enumerate}
 119
 120 The \tr{body} of the original function may not be given (i.e., it's
 121 BOTTOM), in which case you'd jolly well better not tug on the
 122 worker-body output!
 123
 124 Here's an example.  The original function is:
 125 \begin{verbatim}
 126 g :: forall a . Int -> [a] -> a
 127
 128 g = /\ a -> \ x ys ->
 129         case x of
 130           0 -> head ys
 131           _ -> head (tail ys)
 132 \end{verbatim}
 133
 134 From this, we want to produce:
 135 \begin{verbatim}
 136 -- wrapper (an unfolding)
 137 g :: forall a . Int -> [a] -> a
 138
 139 g = /\ a -> \ x ys ->
 140         case x of
 141           I# x# -> g.wrk a x# ys
 142             -- call the worker; don't forget the type args!
 143
 144 -- worker
 145 g.wrk :: forall a . Int# -> [a] -> a
 146
 147 g.wrk = /\ a -> \ x# ys ->
 148         let
 149             x = I# x#
 150         in
 151             case x of               -- note: body of g moved intact
 152               0 -> head ys
 153               _ -> head (tail ys)
 154 \end{verbatim}
 155
 156 Something we have to be careful about:  Here's an example:
 157 \begin{verbatim}
 158 -- "f" strictness: U(P)U(P)
 159 f (I# a) (I# b) = a +# b
 160
 161 g = f   -- "g" strictness same as "f"
 162 \end{verbatim}
 163 \tr{f} will get a worker all nice and friendly-like; that's good.
 164 {\em But we don't want a worker for \tr{g}}, even though it has the
 165 same strictness as \tr{f}.  Doing so could break laziness, at best.
 166
 167 Consequently, we insist that the number of strictness-info items is
 168 exactly the same as the number of lambda-bound arguments.  (This is
 169 probably slightly paranoid, but OK in practice.)  If it isn't the
 170 same, we ``revise'' the strictness info, so that we won't propagate
 171 the unusable strictness-info into the interfaces.
 172
 173 ==========================
 174
 175 Here's the real fun... The wrapper's ``deconstructing'' of arguments
 176 and the worker's putting them back together again are ``duals'' in
 177 some sense.
 178
 179 What we do is walk along the @Demand@ list, producing two
 180 expressions (one for wrapper, one for worker...), each with a ``hole''
 181 in it, where we will later plug in more information.  For our previous
 182 example, the expressions-with-HOLES are:
 183 \begin{verbatim}
 184 \ x ys ->               -- wrapper
 185         case x of
 186           I# x# -> <<HOLE>> x# ys
 187
 188 \ x# ys ->              -- worker
 189         let
 190             x = I# x#
 191         in
 192             <<HOLE>>
 193 \end{verbatim}
 194 (Actually, we add the lambda-bound arguments at the end...) (The big
 195 Lambdas are added on the front later.)
 196
 197 \begin{code}
 198 mkWwBodies
 199         :: Type         -- Type of the *body* of the orig
 200                                 -- function; i.e. /\ tyvars -> \ vars -> body
 201         -> [TyVar]              -- Type lambda vars of original function
 202         -> [Id]                 -- Args of original function
 203         -> [Demand]             -- Strictness info for those args
 204
 205         -> UniqSM (Maybe        -- Nothing iff (a) no interesting split possible
 206                                 --             (b) any unpack on abstract type
 207                      (Id -> CoreExpr,           -- Wrapper expr w/
 208                                                         --   hole for worker id
 209                       CoreExpr -> CoreExpr,     -- Worker expr w/ hole
 210                                                         --   for original fn body
 211                       StrictnessInfo,                   -- Worker strictness info
 212                       Type -> Type)             -- Worker type w/ hole
 213            )                                            --   for type of original fn body
 214
 215
 216 mkWwBodies body_ty tyvars args arg_infos
 217   = ASSERT(length args == length arg_infos)
 218     -- or you can get disastrous user/definer-module mismatches
 219     if (all_absent_args_and_unboxed_value body_ty arg_infos)
 220     then returnUs Nothing
 221
 222     else -- the rest...
 223     mk_ww_arg_processing args arg_infos (mAX_WORKER_ARGS - nonAbsentArgs arg_infos)
 224                 `thenUsMaybe` \ (wrap_frag, work_args_info, work_frag) ->
 225     let
 226         (work_args, wrkr_demands) = unzip work_args_info
 227
 228         wrkr_strictness = mkStrictnessInfo wrkr_demands Nothing -- no worker-of-worker...
 229
 230         wrapper_w_hole = \ worker_id ->
 231                                 mkLam tyvars args (
 232                                 wrap_frag (
 233                                 mkCoTyApps (Var worker_id) (map mkTyVarTy tyvars)
 234                          ))
 235
 236         worker_w_hole = \ orig_body ->
 237                                 mkLam tyvars work_args (
 238                                 work_frag orig_body
 239                         )
 240
 241         worker_ty_w_hole = \ body_ty ->
 242                                 snd (quantifyTy tyvars (
 243                                 mkFunTys (map idType work_args) body_ty
 244                            ))
 245     in
 246     returnUs (Just (wrapper_w_hole, worker_w_hole, wrkr_strictness, worker_ty_w_hole))
 247   where
 248     -- "all_absent_args_and_unboxed_value":
 249     -- check for the obscure case of "\ x y z ... -> body" where
 250     -- (a) *all* of the args x, y, z,... are absent, and
 251     -- (b) the type of body is unboxed
 252     -- If these conditions are true, we must *not* play worker/wrapper games!
 253
 254     all_absent_args_and_unboxed_value body_ty arg_infos
 255       = not (null arg_infos)
 256         && all is_absent_arg arg_infos
 257         && isPrimType body_ty
 258
 259     is_absent_arg (WwLazy True) = True
 260     is_absent_arg _             = False
 261 \end{code}
 262
 263 Important: mk_ww_arg_processing doesn't check
 264 for an "interesting" split.  It just races ahead and makes the
 265 split, even if there's no unpacking at all.  This is important for
 266 when it calls itself recursively.
 267
 268 It returns Nothing only if it encounters an abstract type in mid-flight.
 269
 270 \begin{code}
 271 mAX_WORKER_ARGS :: Int          -- ToDo: set via flag
 272 mAX_WORKER_ARGS = 6             -- Hmm... but this is an everything-must-
 273                                 -- be-compiled-with-the-same-val thing...
 274
 275 mk_ww_arg_processing
 276         :: [Id]                 -- Args of original function
 277         -> [Demand]             -- Strictness info for those args
 278                                 --   must be at least as long as args
 279
 280         -> Int                  -- Number of extra args we are prepared to add.
 281                                 -- This prevents over-eager unpacking, leading
 282                                 -- to huge-arity functions.
 283
 284         -> UniqSM (Maybe        -- Nothing iff any unpack on abstract type
 285                      (CoreExpr -> CoreExpr,     -- Wrapper expr w/
 286                                                         --   hole for worker id
 287                                                         --   applied to types
 288                       [(Id,Demand)],                    -- Worker's args
 289                                                         -- and their strictness info
 290                       CoreExpr -> CoreExpr)     -- Worker body expr w/ hole
 291            )                                            --   for original fn body
 292
 293 mk_ww_arg_processing [] _ _ = returnUs (Just (id, [], id))
 294
 295 mk_ww_arg_processing (arg : args) (WwLazy True : infos) max_extra_args
 296   =     -- Absent argument
 297         -- So, finish args to the right...
 298     --pprTrace "Absent; num_wrkr_args=" (ppInt num_wrkr_args) (
 299     let
 300         arg_ty = idType arg
 301     in
 302     mk_ww_arg_processing args infos max_extra_args
 303                                     -- we've already discounted for absent args,
 304                                     -- so we don't change max_extra_args
 305                    `thenUsMaybe` \ (wrap_rest, work_args_info, work_rest) ->
 306
 307                 -- wrapper doesn't pass this arg to worker:
 308     returnUs (Just (
 309                  -- wrapper:
 310                  \ hole -> wrap_rest hole,
 311
 312                  -- worker:
 313                  work_args_info, -- NB: no argument added
 314                  \ hole -> mk_absent_let arg arg_ty (work_rest hole)
 315     ))
 316     --)
 317   where
 318     mk_absent_let arg arg_ty body
 319       = if not (isPrimType arg_ty) then
 320             Let (NonRec arg (mkCoTyApp (Var aBSENT_ERROR_ID) arg_ty)) body
 321         else -- quite horrible
 322             panic "WwLib: haven't done mk_absent_let for primitives yet"
 323
 324
 325 mk_ww_arg_processing (arg : args) (WwUnpack cmpnt_infos : infos) max_extra_args
 326   | new_max_extra_args > 0      -- Check that we are prepared to add arguments
 327   =     -- this is the complicated one.
 328     --pprTrace "Unpack; num_wrkr_args=" (ppCat [ppInt num_wrkr_args, ppStr "; new_max=", ppInt new_num_wrkr_args, ppStr "; arg=", ppr PprDebug arg, ppr PprDebug (WwUnpack cmpnt_infos)]) (
 329     case maybeDataTyCon arg_ty of
 330
 331           Nothing         ->       -- Not a data type
 332                                    panic "mk_ww_arg_processing: not datatype"
 333
 334           Just (_, _, []) ->       -- An abstract type
 335                                    -- We have to give up on the whole idea
 336                                    returnUs Nothing
 337           Just (_, _, (_:_:_)) ->  -- Two or more constructors; that's odd
 338                                    panic "mk_ww_arg_processing: multi-constr"
 339
 340           Just (arg_tycon, tycon_arg_tys, [data_con]) ->
 341                         -- The main event: a single-constructor data type
 342
 343             let
 344                 (_,inst_con_arg_tys,_)
 345                   = getInstantiatedDataConSig data_con tycon_arg_tys
 346             in
 347             getUniques (length inst_con_arg_tys)    `thenUs` \ uniqs ->
 348
 349             let
 350                 unpk_args = zipWithEqual
 351                              (\ u t -> mkSysLocal SLIT("upk") u t mkUnknownSrcLoc)
 352                              uniqs inst_con_arg_tys
 353             in
 354                 -- In processing the rest, push the sub-component args
 355                 -- and infos on the front of the current bunch
 356             mk_ww_arg_processing (unpk_args ++ args) (cmpnt_infos ++ infos) new_max_extra_args
 357                         `thenUsMaybe` \ (wrap_rest, work_args_info, work_rest) ->
 358
 359             returnUs (Just (
 360               -- wrapper: unpack the value
 361               \ hole -> mk_unpk_case arg unpk_args
 362                             data_con arg_tycon
 363                             (wrap_rest hole),
 364
 365               -- worker: expect the unpacked value;
 366               -- reconstruct the orig value with a "let"
 367               work_args_info,
 368               \ hole -> work_rest (mk_pk_let arg data_con tycon_arg_tys unpk_args hole)
 369             ))
 370     --)
 371   where
 372     arg_ty = idType arg
 373
 374     new_max_extra_args
 375       = max_extra_args
 376         + 1                         -- We won't pass the original arg now
 377         - nonAbsentArgs cmpnt_infos -- But we will pass an arg for each cmpt
 378
 379     mk_unpk_case arg unpk_args boxing_con boxing_tycon body
 380       = Case (Var arg) (
 381           AlgAlts [(boxing_con, unpk_args, body)]
 382           NoDefault
 383         )
 384
 385     mk_pk_let arg boxing_con con_tys unpk_args body
 386       = Let (NonRec arg (Con boxing_con con_tys [VarArg a | a <- unpk_args]))
 387               body
 388
 389 mk_ww_arg_processing (arg : args) (arg_demand : infos) max_extra_args
 390   | otherwise
 391   =     -- For all others at the moment, we just
 392         -- pass them to the worker unchanged.
 393     --pprTrace "Other; num_wrkr_args=" (ppCat [ppInt num_wrkr_args, ppStr ";arg=", ppr PprDebug arg, ppr PprDebug arg_demand]) (
 394
 395         -- Finish args to the right...
 396     mk_ww_arg_processing args infos max_extra_args
 397                         `thenUsMaybe` \ (wrap_rest, work_args_info, work_rest) ->
 398
 399     returnUs (Just (
 400               -- wrapper:
 401               \ hole -> wrap_rest (App hole (VarArg arg)),
 402
 403               -- worker:
 404               (arg, arg_demand) : work_args_info,
 405               \ hole -> work_rest hole
 406     ))
 407     --)
 408 \end{code}
 409
 410 %************************************************************************
 411 %*                                                                      *
 412 \subsection[monad-WwLib]{Simple monad for worker/wrapper}
 413 %*                                                                      *
 414 %************************************************************************
 415
 416 In this monad, we thread a @UniqueSupply@, and we carry a
 417 @GlobalSwitch@-lookup function downwards.
 418
 419 \begin{code}
 420 type WwM result
 421   =  UniqSupply
 422   -> (GlobalSwitch -> Bool)
 423   -> result
 424
 425 {-# INLINE thenWw #-}
 426 {-# INLINE returnWw #-}
 427
 428 returnWw :: a -> WwM a
 429 thenWw   :: WwM a -> (a -> WwM b) -> WwM b
 430 mapWw    :: (a -> WwM b) -> [a] -> WwM [b]
 431
 432 returnWw expr ns sw = expr
 433
 434 thenWw m k us sw_chk
 435   = case splitUniqSupply us     of { (s1, s2) ->
 436     case (m s1 sw_chk)          of { m_res ->
 437     k m_res s2 sw_chk }}
 438
 439 mapWw f []     = returnWw []
 440 mapWw f (x:xs)
 441   = f x         `thenWw` \ x'  ->
 442     mapWw f xs  `thenWw` \ xs' ->
 443     returnWw (x':xs')
 444 \end{code}
 445
 446 \begin{code}
 447 getUniqueWw :: WwM Unique
 448 uniqSMtoWwM :: UniqSM a -> WwM a
 449
 450 getUniqueWw us sw_chk = getUnique us
 451
 452 uniqSMtoWwM u_obj us sw_chk = u_obj us
 453
 454 thenUsMaybe :: UniqSM (Maybe a) -> (a -> UniqSM (Maybe b)) -> UniqSM (Maybe b)
 455 thenUsMaybe m k
 456   = m   `thenUs` \ result ->
 457     case result of
 458       Nothing -> returnUs Nothing
 459       Just x  -> k x
 460 \end{code}