ghc/compiler/simplCore/FoldrBuildWW.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
   3 %
   4 \section[FoldrBuildWW]{Spliting suitable functions into Workers and Wrappers}
   5
   6 \begin{code}
   7 #include "HsVersions.h"
   8
   9 module FoldrBuildWW ( mkFoldrBuildWW ) where
  10
  11 IMPORT_Trace
  12 import Outputable
  13 import Pretty
  14 import Type             ( cloneTyVarFromTemplate, mkTyVarTy,
  15                           splitTypeWithDictsAsArgs, eqTyCon,  mkForallTy )
  16 import TysPrim          ( alphaTy )
  17 import TyVar            ( alphaTyVar )
  18
  19 import Type             ( Type(..) ) -- **** CAN SEE THE CONSTRUCTORS ****
  20 import UniqSupply       ( runBuiltinUs )
  21 import WwLib            -- share the same monad (is this eticit ?)
  22 import PrelInfo         ( listTyCon, mkListTy, nilDataCon, consDataCon,
  23                           foldrId, buildId
  24                         )
  25 import Id               ( getIdFBTypeInfo, mkWorkerId, getIdInfo,
  26                           replaceIdInfo, mkSysLocal, idType
  27                         )
  28 import IdInfo
  29 import Maybes
  30 import SrcLoc           ( mkUnknownSrcLoc, SrcLoc )
  31 import Util
  32 \end{code}
  33
  34 \begin{code}
  35 mkFoldrBuildWW
  36         :: (GlobalSwitch -> Bool)
  37         -> UniqSupply
  38         -> [CoreBinding]
  39         -> [CoreBinding]
  40 mkFoldrBuildWW switch us top_binds =
  41    (mapWw wwBind top_binds `thenWw` \ top_binds2 ->
  42    returnWw (concat top_binds2)) us switch
  43 \end{code}
  44
  45 \begin{code}
  46 wwBind :: CoreBinding -> WwM [CoreBinding]
  47 wwBind (NonRec bndr expr)
  48   = try_split_bind bndr expr    `thenWw` \ re ->
  49     returnWw [NonRec bnds expr | (bnds,expr) <- re]
  50 wwBind (Rec binds)
  51   = mapWw (\ (bndr,expr) -> try_split_bind bndr expr) binds   `thenWw` \ res ->
  52     returnWw [Rec (concat res)]
  53
  54 wwExpr :: CoreExpr -> WwM CoreExpr
  55 wwExpr e@(Var _) = returnWw e
  56 wwExpr e@(Lit _) = returnWw e
  57 wwExpr e@(Con _ _ _) = returnWw e
  58 wwExpr e@(Prim _ _ _) = returnWw e
  59 wwExpr   (Lam ids e) =
  60         wwExpr e                `thenWw` \ e' ->
  61         returnWw (Lam ids e')
  62 wwExpr   (CoTyLam tyvar e) =
  63         wwExpr e                `thenWw` \ e' ->
  64         returnWw (CoTyLam tyvar e')
  65 wwExpr   (App f atom) =
  66         wwExpr f                `thenWw` \ f' ->
  67         returnWw (App f atom)
  68 wwExpr   (CoTyApp f ty) =
  69         wwExpr f                `thenWw` \ f' ->
  70         returnWw (CoTyApp f' ty)
  71 wwExpr   (SCC lab e) =
  72         wwExpr e                `thenWw` \ e' ->
  73         returnWw (SCC lab e')
  74 wwExpr   (Let bnds e) =
  75         wwExpr e                `thenWw` \ e' ->
  76         wwBind bnds             `thenWw` \ bnds' ->
  77         returnWw (foldr Let e' bnds')
  78 wwExpr   (Case e alts) =
  79         wwExpr e                `thenWw` \ e' ->
  80         wwAlts alts             `thenWw` \ alts' ->
  81         returnWw  (Case e' alts')
  82
  83 wwAlts (AlgAlts alts deflt) =
  84         mapWw (\(con,binders,e) ->
  85                         wwExpr e        `thenWw` \ e' ->
  86                         returnWw (con,binders,e')) alts `thenWw` \ alts' ->
  87         wwDef deflt                                     `thenWw` \ deflt' ->
  88         returnWw (AlgAlts alts' deflt)
  89 wwAlts (PrimAlts alts deflt) =
  90         mapWw (\(lit,e) ->
  91                         wwExpr e        `thenWw` \ e' ->
  92                         returnWw (lit,e')) alts         `thenWw` \ alts' ->
  93         wwDef deflt                                     `thenWw` \ deflt' ->
  94         returnWw (PrimAlts alts' deflt)
  95
  96 wwDef e@NoDefault = returnWw e
  97 wwDef  (BindDefault v e) =
  98         wwExpr e                                        `thenWw` \ e' ->
  99         returnWw (BindDefault v e')
 100 \end{code}
 101
 102 \begin{code}
 103 try_split_bind :: Id -> CoreExpr -> WwM [(Id,CoreExpr)]
 104 try_split_bind id expr =
 105   wwExpr expr                   `thenWw` \ expr' ->
 106   case getFBType (getIdFBTypeInfo id) of
 107     Just (FBType consum prod)
 108         |  FBGoodProd == prod ->
 109 {-      || any (== FBGoodConsum) consum -}
 110       let
 111         (use_args,big_args,args,body) = digForLambdas expr'
 112       in
 113         if length args /= length consum   -- funny number of arguments
 114         then returnWw [(id,expr')]
 115         else
 116         -- f /\ t1 .. tn \ v1 .. vn -> e
 117         --      ===>
 118         -- f_wrk /\ t1 .. tn t_new \ v1 .. vn c n -> foldr <exprTy> <nTy> c n e
 119         -- f /\ t1 .. tn \ v1 .. vn
 120         --      -> build exprTy (\ c n -> f_wrk t1 .. tn t_new v1 .. vn c n)
 121         pprTrace "WW:" (ppr PprDebug id) (returnWw ())
 122                                 `thenWw` \ () ->
 123         getUniqueWw             `thenWw` \ ty_new_uq ->
 124         getUniqueWw             `thenWw` \ worker_new_uq ->
 125         getUniqueWw             `thenWw` \ c_new_uq ->
 126         getUniqueWw             `thenWw` \ n_new_uq ->
 127       let
 128         -- The *new* type
 129         n_ty = alphaTy
 130         n_ty_templ = alphaTy
 131
 132         (templ,arg_tys,res) = splitTypeWithDictsAsArgs (idType id)
 133         expr_ty = getListTy res
 134         getListTy res = panic "FoldrBuildWW:getListTy:ToDo" {-LATER:case res of
 135                          UniData lty [ty] | lty `eqTyCon` listTyCon -> ty
 136                          _ -> panic "Trying to split a non List datatype into Worker/Wrapper"-}
 137
 138         c_ty       = expr_ty `mkFunTy` (n_ty `mkFunTy` n_ty)
 139         c_ty_templ = expr_ty `mkFunTy` (n_ty_templ `mkFunTy` n_ty_templ)
 140
 141         worker_ty = mkForallTy (templ  ++ [alphaTyVar])
 142                         (foldr mkFunTy n_ty_templ (arg_tys++[c_ty_templ,n_ty_templ]))
 143         wrapper_id  = id `replaceIdInfo`
 144                               (getIdInfo id     `addInfo_UF`
 145                                iWantToBeINLINEd UnfoldAlways)
 146         worker_id  = mkWorkerId worker_new_uq id worker_ty
 147                                 noIdInfo
 148                 -- TODO : CHECK if mkWorkerId is thr
 149                 -- right function to use ..
 150         -- Now the bodies
 151
 152         c_id = mkSysLocal SLIT("_fbww") c_new_uq c_ty mkUnknownSrcLoc
 153         n_id = mkSysLocal SLIT("_fbww") n_new_uq n_ty mkUnknownSrcLoc
 154         worker_rhs
 155           = mkTyLam [] (big_args ++ [alphaTyVar]) (args++[c_id,n_id]) worker_body
 156
 157         worker_body = runBuiltinUs (
 158           mkCoApps
 159             (Var foldrId `CoTyApp` expr_ty `CoTyApp` n_ty `App`
 160                VarArg c_id `App` VarArg n_id)
 161             [body])
 162         wrapper_rhs = mkLam big_args args wrapper_body
 163
 164         wrapper_body = runBuiltinUs (
 165                  mkCoApps (CoTyApp (Var buildId) expr_ty)
 166                                 [mkLam [alphaTyVar] [c_id,n_id]
 167                 (foldl App
 168                         (mkCoTyApps (Var worker_id)
 169                                 [mkTyVarTy t | t <- big_args ++ [alphaTyVar]])
 170                         (map VarArg (args++[c_id,n_id])))])
 171
 172       in
 173         if length args /= length arg_tys ||
 174            length big_args /= length templ
 175         then panic "LEN PROBLEM"
 176         else
 177         returnWw  [(worker_id,worker_rhs),(wrapper_id,wrapper_rhs)]
 178     _ -> returnWw [(id,expr')]
 179 \end{code}
 180