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         mkWwBodies,
   9         worthSplitting, setUnpackStrategy
  10     ) where
  11
  12 #include "HsVersions.h"
  13
  14 import CoreSyn
  15 import CoreUtils        ( exprType )
  16 import Id               ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
  17                           isOneShotLambda, setOneShotLambda,
  18                           setIdInfo
  19                         )
  20 import IdInfo           ( CprInfo(..), vanillaIdInfo )
  21 import DataCon          ( splitProductType )
  22 import Demand           ( Demand(..), wwLazy, wwPrim )
  23 import PrelInfo         ( realWorldPrimId, aBSENT_ERROR_ID )
  24 import TysPrim          ( realWorldStatePrimTy )
  25 import TysWiredIn       ( tupleCon )
  26 import Type             ( Type, isUnLiftedType, mkFunTys,
  27                           splitForAllTys, splitFunTys,  isAlgType
  28                         )
  29 import BasicTypes       ( Arity, Boxity(..) )
  30 import Var              ( Var, isId )
  31 import UniqSupply       ( returnUs, thenUs, getUniqueUs, getUniquesUs, UniqSM )
  32 import Util             ( zipWithEqual )
  33 import Outputable
  34 import List             ( zipWith4 )
  35 \end{code}
  36
  37
  38 %************************************************************************
  39 %*                                                                      *
  40 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
  41 %*                                                                      *
  42 %************************************************************************
  43
  44         ************   WARNING  ******************
  45         these comments are rather out of date
  46         *****************************************
  47
  48 @mkWrapperAndWorker@ is given:
  49 \begin{enumerate}
  50 \item
  51 The {\em original function} \tr{f}, of the form:
  52 \begin{verbatim}
  53 f = /\ tyvars -> \ args -> body
  54 \end{verbatim}
  55 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
  56 are given separately.
  57
  58 We use the Id \tr{f} mostly to get its type.
  59
  60 \item
  61 Strictness information about \tr{f}, in the form of a list of
  62 @Demands@.
  63
  64 \item
  65 A @UniqueSupply@.
  66 \end{enumerate}
  67
  68 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
  69 \begin{enumerate}
  70 \item
  71 Maybe @Nothing@: no worker/wrappering going on in this case. This can
  72 happen (a)~if the strictness info says that there is nothing
  73 interesting to do or (b)~if *any* of the argument types corresponding
  74 to ``active'' arg postitions is abstract or will be to the outside
  75 world (i.e., {\em this} module can see the constructors, but nobody
  76 else will be able to).  An ``active'' arg position is one which the
  77 wrapper has to unpack.  An importing module can't do this unpacking,
  78 so it simply has to give up and call the wrapper only.
  79
  80 \item
  81 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
  82
  83 The @wrapper_Id@ is just the one that was passed in, with its
  84 strictness IdInfo updated.
  85 \end{enumerate}
  86
  87 The \tr{body} of the original function may not be given (i.e., it's
  88 BOTTOM), in which case you'd jolly well better not tug on the
  89 worker-body output!
  90
  91 Here's an example.  The original function is:
  92 \begin{verbatim}
  93 g :: forall a . Int -> [a] -> a
  94
  95 g = /\ a -> \ x ys ->
  96         case x of
  97           0 -> head ys
  98           _ -> head (tail ys)
  99 \end{verbatim}
 100
 101 From this, we want to produce:
 102 \begin{verbatim}
 103 -- wrapper (an unfolding)
 104 g :: forall a . Int -> [a] -> a
 105
 106 g = /\ a -> \ x ys ->
 107         case x of
 108           I# x# -> g.wrk a x# ys
 109             -- call the worker; don't forget the type args!
 110
 111 -- worker
 112 g.wrk :: forall a . Int# -> [a] -> a
 113
 114 g.wrk = /\ a -> \ x# ys ->
 115         let
 116             x = I# x#
 117         in
 118             case x of               -- note: body of g moved intact
 119               0 -> head ys
 120               _ -> head (tail ys)
 121 \end{verbatim}
 122
 123 Something we have to be careful about:  Here's an example:
 124 \begin{verbatim}
 125 -- "f" strictness: U(P)U(P)
 126 f (I# a) (I# b) = a +# b
 127
 128 g = f   -- "g" strictness same as "f"
 129 \end{verbatim}
 130 \tr{f} will get a worker all nice and friendly-like; that's good.
 131 {\em But we don't want a worker for \tr{g}}, even though it has the
 132 same strictness as \tr{f}.  Doing so could break laziness, at best.
 133
 134 Consequently, we insist that the number of strictness-info items is
 135 exactly the same as the number of lambda-bound arguments.  (This is
 136 probably slightly paranoid, but OK in practice.)  If it isn't the
 137 same, we ``revise'' the strictness info, so that we won't propagate
 138 the unusable strictness-info into the interfaces.
 139
 140
 141 %************************************************************************
 142 %*                                                                      *
 143 \subsection{Functions over Demands}
 144 %*                                                                      *
 145 %************************************************************************
 146
 147 \begin{code}
 148 mAX_WORKER_ARGS :: Int          -- ToDo: set via flag
 149 mAX_WORKER_ARGS = 6
 150
 151 setUnpackStrategy :: [Demand] -> [Demand]
 152 setUnpackStrategy ds
 153   = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
 154   where
 155     go :: Int                   -- Max number of args available for sub-components of [Demand]
 156        -> [Demand]
 157        -> (Int, [Demand])       -- Args remaining after subcomponents of [Demand] are unpacked
 158
 159     go n (WwUnpack _ cs : ds) | n' >= 0
 160                               = WwUnpack True cs' `cons` go n'' ds
 161                               | otherwise
 162                               = WwUnpack False cs `cons` go n ds
 163                                  where
 164                                    n' = n + 1 - nonAbsentArgs cs
 165                                         -- Add one because we don't pass the top-level arg any more
 166                                         -- Delete # of non-absent args to which we'll now be committed
 167                                    (n'',cs') = go n' cs
 168
 169     go n (d:ds) = d `cons` go n ds
 170     go n []     = (n,[])
 171
 172     cons d (n,ds) = (n, d:ds)
 173
 174 nonAbsentArgs :: [Demand] -> Int
 175 nonAbsentArgs []                 = 0
 176 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
 177 nonAbsentArgs (d           : ds) = 1 + nonAbsentArgs ds
 178
 179 worthSplitting :: [Demand]
 180                -> Bool  -- Result is bottom
 181                -> Bool  -- True <=> the wrapper would not be an identity function
 182 worthSplitting ds result_bot = any worth_it ds
 183         -- We used not to split if the result is bottom.
 184         -- [Justification:  there's no efficiency to be gained.]
 185         -- But it's sometimes bad not to make a wrapper.  Consider
 186         --      fw = \x# -> let x = I# x# in case e of
 187         --                                      p1 -> error_fn x
 188         --                                      p2 -> error_fn x
 189         --                                      p3 -> the real stuff
 190         -- The re-boxing code won't go away unless error_fn gets a wrapper too.
 191
 192   where
 193     worth_it (WwLazy True)     = True   -- Absent arg
 194     worth_it (WwUnpack True _) = True   -- Arg to unpack
 195     worth_it WwStrict          = False  -- Don't w/w just because of strictness
 196     worth_it other             = False
 197
 198 allAbsent :: [Demand] -> Bool
 199 allAbsent ds = all absent ds
 200   where
 201     absent (WwLazy is_absent) = is_absent
 202     absent (WwUnpack True cs) = allAbsent cs
 203     absent other              = False
 204 \end{code}
 205
 206
 207 %************************************************************************
 208 %*                                                                      *
 209 \subsection{The worker wrapper core}
 210 %*                                                                      *
 211 %************************************************************************
 212
 213 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
 214
 215 \begin{code}
 216 mkWwBodies :: Type                              -- Type of original function
 217            -> Arity                             -- Arity of original function
 218            -> [Demand]                          -- Strictness of original function
 219            -> Bool                              -- True <=> function returns bottom
 220            -> [Bool]                            -- One-shot-ness of the function
 221            -> CprInfo                           -- Result of CPR analysis
 222            -> UniqSM ([Demand],                 -- Demands for worker (value) args
 223                       Id -> CoreExpr,           -- Wrapper body, lacking only the worker Id
 224                       CoreExpr -> CoreExpr)     -- Worker body, lacking the original function rhs
 225
 226 -- wrap_fn_args E       = \x y -> E
 227 -- work_fn_args E       = E x y
 228
 229 -- wrap_fn_str E        = case x of { (a,b) ->
 230 --                        case a of { (a1,a2) ->
 231 --                        E a1 a2 b y }}
 232 -- work_fn_str E        = \a2 a2 b y ->
 233 --                        let a = (a1,a2) in
 234 --                        let x = (a,b) in
 235 --                        E
 236
 237 mkWwBodies fun_ty arity demands res_bot one_shots cpr_info
 238   = mkWWargs fun_ty arity demands' res_bot one_shots'   `thenUs` \ (wrap_args, wrap_fn_args,   work_fn_args, res_ty) ->
 239     mkWWcpr res_ty cpr_info                             `thenUs` \ (wrap_fn_cpr,    work_fn_cpr,  cpr_res_ty) ->
 240     mkWWstr cpr_res_ty wrap_args                        `thenUs` \ (work_dmds, wrap_fn_str,    work_fn_str) ->
 241
 242     returnUs (work_dmds,
 243               Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . Var,
 244               work_fn_str . work_fn_cpr . work_fn_args)
 245         -- We use an INLINE unconditionally, even if the wrapper turns out to be
 246         -- something trivial like
 247         --      fw = ...
 248         --      f = __inline__ (coerce T fw)
 249         -- The point is to propagate the coerce to f's call sites, so even though
 250         -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
 251         -- fw from being inlined into f's RHS
 252   where
 253     demands'   = demands   ++ repeat wwLazy
 254     one_shots' = one_shots ++ repeat False
 255 \end{code}
 256
 257
 258 %************************************************************************
 259 %*                                                                      *
 260 \subsection{Coercion stuff}
 261 %*                                                                      *
 262 %************************************************************************
 263
 264
 265 We really want to "look through" coerces.
 266 Reason: I've seen this situation:
 267
 268         let f = coerce T (\s -> E)
 269         in \x -> case x of
 270                     p -> coerce T' f
 271                     q -> \s -> E2
 272                     r -> coerce T' f
 273
 274 If only we w/w'd f, we'd get
 275         let f = coerce T (\s -> fw s)
 276             fw = \s -> E
 277         in ...
 278
 279 Now we'll inline f to get
 280
 281         let fw = \s -> E
 282         in \x -> case x of
 283                     p -> fw
 284                     q -> \s -> E2
 285                     r -> fw
 286
 287 Now we'll see that fw has arity 1, and will arity expand
 288 the \x to get what we want.
 289
 290 \begin{code}
 291 -- mkWWargs is driven off the function type and arity.
 292 -- It chomps bites off foralls, arrows, newtypes
 293 -- and keeps repeating that until it's satisfied the supplied arity
 294
 295 mkWWargs :: Type -> Arity
 296          -> [Demand] -> Bool -> [Bool]          -- Both these will in due course be derived
 297                                                 -- from the type.  The [Bool] is True for a one-shot arg.
 298                                                 -- ** Both are infinite, extended with neutral values if necy **
 299          -> UniqSM  ([Var],             -- Wrapper args
 300                      CoreExpr -> CoreExpr,      -- Wrapper fn
 301                      CoreExpr -> CoreExpr,      -- Worker fn
 302                      Type)                      -- Type of wrapper body
 303
 304 mkWWargs fun_ty arity demands res_bot one_shots
 305   | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
 306         -- If the function returns bottom, we feel free to
 307         -- build lots of wrapper args:
 308         --        \x. let v=E in \y. bottom
 309         --      = \xy. let v=E in bottom
 310   = getUniquesUs                `thenUs` \ wrap_uniqs ->
 311     let
 312       val_args  = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
 313       wrap_args = tyvars ++ val_args
 314     in
 315     mkWWargs new_fun_ty
 316              (arity - n_args)
 317              (drop n_args demands)
 318              res_bot
 319              (drop n_args one_shots)    `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
 320
 321     returnUs (wrap_args ++ more_wrap_args,
 322               mkLams wrap_args . wrap_fn_args,
 323               work_fn_args . applyToVars wrap_args,
 324               res_ty)
 325   where
 326     (tyvars, tau)       = splitForAllTys fun_ty
 327     (arg_tys, body_ty)  = splitFunTys tau
 328     n_arg_tys           = length arg_tys
 329     n_args              | res_bot   = n_arg_tys
 330                         | otherwise = arity `min` n_arg_tys
 331     new_fun_ty          | n_args == n_arg_tys = body_ty
 332                         | otherwise           = mkFunTys (drop n_args arg_tys) body_ty
 333
 334 mkWWargs fun_ty arity demands res_bot one_shots
 335   = returnUs ([], id, id, fun_ty)
 336
 337 applyToVars :: [Var] -> CoreExpr -> CoreExpr
 338 applyToVars vars fn = mkVarApps fn vars
 339
 340 mk_wrap_arg uniq ty dmd one_shot
 341   = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
 342   where
 343     set_one_shot True  id = setOneShotLambda id
 344     set_one_shot False id = id
 345 \end{code}
 346
 347
 348 %************************************************************************
 349 %*                                                                      *
 350 \subsection{Strictness stuff}
 351 %*                                                                      *
 352 %************************************************************************
 353
 354 \begin{code}
 355 mkWWstr :: Type                                 -- Result type
 356         -> [Var]                                -- Wrapper args; have their demand info on them
 357                                                 -- *Includes type variables*
 358         -> UniqSM ([Demand],                    -- Demand on worker (value) args
 359                    CoreExpr -> CoreExpr,        -- Wrapper body, lacking the worker call
 360                                                 -- and without its lambdas
 361                                                 -- This fn adds the unboxing, and makes the
 362                                                 -- call passing the unboxed things
 363
 364                    CoreExpr -> CoreExpr)        -- Worker body, lacking the original body of the function,
 365                                                 -- but *with* lambdas
 366
 367 mkWWstr res_ty wrap_args
 368   = mk_ww_str wrap_args         `thenUs` \ (work_args, take_apart, put_together) ->
 369     let
 370         work_dmds = [idDemandInfo v | v <- work_args, isId v]
 371         apply_to args fn = mkVarApps fn args
 372     in
 373     if not (null work_dmds && isUnLiftedType res_ty) then
 374         returnUs ( work_dmds,
 375                    take_apart . apply_to work_args,
 376                    mkLams work_args . put_together)
 377     else
 378         -- Horrid special case.  If the worker would have no arguments, and the
 379         -- function returns a primitive type value, that would make the worker into
 380         -- an unboxed value.  We box it by passing a dummy void argument, thus:
 381         --
 382         --      f = /\abc. \xyz. fw abc void
 383         --      fw = /\abc. \v. body
 384         --
 385         -- We use the state-token type which generates no code
 386     getUniqueUs                 `thenUs` \ void_arg_uniq ->
 387     let
 388         void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
 389     in
 390     returnUs ([wwPrim],
 391               take_apart . apply_to [realWorldPrimId] . apply_to work_args,
 392               mkLams work_args . Lam void_arg . put_together)
 393
 394         -- Empty case
 395 mk_ww_str []
 396   = returnUs ([],
 397               \ wrapper_body -> wrapper_body,
 398               \ worker_body  -> worker_body)
 399
 400
 401 mk_ww_str (arg : ds)
 402   | isTyVar arg
 403   = mk_ww_str ds                `thenUs` \ (worker_args, wrap_fn, work_fn) ->
 404     returnUs (arg : worker_args, wrap_fn, work_fn)
 405
 406   | otherwise
 407   = case idDemandInfo arg of
 408
 409         -- Absent case
 410       WwLazy True ->
 411         mk_ww_str ds            `thenUs` \ (worker_args, wrap_fn, work_fn) ->
 412         returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
 413
 414         -- Unpack case
 415       WwUnpack True cs ->
 416         getUniquesUs            `thenUs` \ uniqs ->
 417         let
 418           unpk_args      = zipWith mk_ww_local uniqs inst_con_arg_tys
 419           unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
 420         in
 421         mk_ww_str (unpk_args_w_ds ++ ds)                `thenUs` \ (worker_args, wrap_fn, work_fn) ->
 422         returnUs (worker_args,
 423                   mk_unpk_case arg unpk_args data_con arg_tycon . wrap_fn,
 424                   work_fn . mk_pk_let arg data_con tycon_arg_tys unpk_args)
 425         where
 426           (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
 427
 428         -- Other cases
 429       other_demand ->
 430         mk_ww_str ds            `thenUs` \ (worker_args, wrap_fn, work_fn) ->
 431         returnUs (arg : worker_args, wrap_fn, work_fn)
 432   where
 433         -- If the wrapper argument is a one-shot lambda, then
 434         -- so should (all) the corresponding worker arguments be
 435         -- This bites when we do w/w on a case join point
 436     set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
 437
 438     set_one_shot | isOneShotLambda arg = setOneShotLambda
 439                  | otherwise           = \x -> x
 440 \end{code}
 441
 442
 443 %************************************************************************
 444 %*                                                                      *
 445 \subsection{CPR stuff}
 446 %*                                                                      *
 447 %************************************************************************
 448
 449
 450 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
 451 info and adds in the CPR transformation.  The worker returns an
 452 unboxed tuple containing non-CPR components.  The wrapper takes this
 453 tuple and re-produces the correct structured output.
 454
 455 The non-CPR results appear ordered in the unboxed tuple as if by a
 456 left-to-right traversal of the result structure.
 457
 458
 459 \begin{code}
 460 mkWWcpr :: Type                              -- function body type
 461         -> CprInfo                           -- CPR analysis results
 462         -> UniqSM (CoreExpr -> CoreExpr,             -- New wrapper
 463                    CoreExpr -> CoreExpr,             -- New worker
 464                    Type)                        -- Type of worker's body
 465
 466 mkWWcpr body_ty NoCPRInfo
 467     = returnUs (id, id, body_ty)      -- Must be just the strictness transf.
 468
 469 mkWWcpr body_ty ReturnsCPR
 470     | not (isAlgType body_ty)
 471     = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
 472       returnUs (id, id, body_ty)
 473
 474     | n_con_args == 1 && isUnLiftedType con_arg_ty1
 475         -- Special case when there is a single result of unlifted type
 476     = getUniquesUs                      `thenUs` \ (work_uniq : arg_uniq : _) ->
 477       let
 478         work_wild = mk_ww_local work_uniq body_ty
 479         arg       = mk_ww_local arg_uniq  con_arg_ty1
 480       in
 481       returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
 482                 \ body     -> workerCase body work_wild [(DataAlt data_con, [arg], Var arg)],
 483                 con_arg_ty1)
 484
 485     | otherwise         -- The general case
 486     = getUniquesUs              `thenUs` \ uniqs ->
 487       let
 488         (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
 489         arg_vars                       = map Var args
 490         ubx_tup_con                    = tupleCon Unboxed n_con_args
 491         ubx_tup_ty                     = exprType ubx_tup_app
 492         ubx_tup_app                    = mkConApp ubx_tup_con (map Type con_arg_tys   ++ arg_vars)
 493         con_app                        = mkConApp data_con    (map Type tycon_arg_tys ++ arg_vars)
 494       in
 495       returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
 496                 \ body     -> workerCase body work_wild [(DataAlt data_con,    args, ubx_tup_app)],
 497                 ubx_tup_ty)
 498     where
 499       (_, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
 500       n_con_args  = length con_arg_tys
 501       con_arg_ty1 = head con_arg_tys
 502
 503 -- If the original function looked like
 504 --      f = \ x -> _scc_ "foo" E
 505 --
 506 -- then we want the CPR'd worker to look like
 507 --      \ x -> _scc_ "foo" (case E of I# x -> x)
 508 -- and definitely not
 509 --      \ x -> case (_scc_ "foo" E) of I# x -> x)
 510 --
 511 -- This transform doesn't move work or allocation
 512 -- from one cost centre to another
 513
 514 workerCase (Note (SCC cc) e) arg alts = Note (SCC cc) (Case e arg alts)
 515 workerCase e                 arg alts = Case e arg alts
 516 \end{code}
 517
 518
 519 %************************************************************************
 520 %*                                                                      *
 521 \subsection{Utilities}
 522 %*                                                                      *
 523 %************************************************************************
 524
 525
 526 \begin{code}
 527 mk_absent_let arg body
 528   | not (isUnLiftedType arg_ty)
 529   = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
 530   | otherwise
 531   = panic "WwLib: haven't done mk_absent_let for primitives yet"
 532   where
 533     arg_ty = idType arg
 534
 535 mk_unpk_case arg unpk_args boxing_con boxing_tycon body
 536         -- A data type
 537   = Case (Var arg)
 538          (sanitiseCaseBndr arg)
 539          [(DataAlt boxing_con, unpk_args, body)]
 540
 541 sanitiseCaseBndr :: Id -> Id
 542 -- The argument we are scrutinising has the right type to be
 543 -- a case binder, so it's convenient to re-use it for that purpose.
 544 -- But we *must* throw away all its IdInfo.  In particular, the argument
 545 -- will have demand info on it, and that demand info may be incorrect for
 546 -- the case binder.  e.g.       case ww_arg of ww_arg { I# x -> ... }
 547 -- Quite likely ww_arg isn't used in '...'.  The case may get discarded
 548 -- if the case binder says "I'm demanded".  This happened in a situation
 549 -- like         (x+y) `seq` ....
 550 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
 551
 552 mk_pk_let arg boxing_con con_tys unpk_args body
 553   = Let (NonRec arg (mkConApp boxing_con con_args)) body
 554   where
 555     con_args = map Type con_tys ++ map Var unpk_args
 556
 557
 558 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty
 559
 560 \end{code}