ghc/compiler/simplCore/FoldrBuildWW.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
   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 AbsUniType       ( alpha_tv, cloneTyVarFromTemplate, mkTyVarTy,
  15                           splitTypeWithDictsAsArgs, eqTyCon,  mkForallTy,
  16                           alpha_tyvar, alpha_ty, alpha, TyVarTemplate
  17                           IF_ATTACK_PRAGMAS(COMMA cmpTyCon)
  18                         )
  19 import UniType          ( UniType(..) ) -- **** CAN SEE THE CONSTRUCTORS ****
  20 import PlainCore
  21 import Unique           ( runBuiltinUs )
  22 import WwLib            -- share the same monad (is this eticit ?)
  23 import AbsPrel          ( listTyCon, mkListTy, nilDataCon, consDataCon,
  24                           foldrId, mkBuild, mkFoldr, buildId,
  25                           mkFunTy
  26                         )
  27 import Id               ( getIdFBTypeInfo, mkWorkerId, getIdInfo,
  28                           replaceIdInfo, mkSysLocal, getIdUniType
  29                         )
  30 import IdInfo
  31 import Maybes
  32 import SrcLoc           ( mkUnknownSrcLoc, SrcLoc )
  33 import Util
  34 \end{code}
  35
  36 \begin{code}
  37 mkFoldrBuildWW
  38         :: (GlobalSwitch -> Bool)
  39         -> SplitUniqSupply
  40         -> PlainCoreProgram
  41         -> PlainCoreProgram
  42 mkFoldrBuildWW switch us top_binds =
  43    (mapWw wwBind top_binds `thenWw` \ top_binds2 ->
  44    returnWw (concat top_binds2)) us switch
  45 \end{code}
  46
  47 \begin{code}
  48 wwBind :: PlainCoreBinding -> WwM [PlainCoreBinding]
  49 wwBind (CoNonRec bndr expr)
  50   = try_split_bind bndr expr    `thenWw` \ re ->
  51     returnWw [CoNonRec bnds expr | (bnds,expr) <- re]
  52 wwBind (CoRec binds)
  53   = mapWw (\ (bndr,expr) -> try_split_bind bndr expr) binds   `thenWw` \ res ->
  54     returnWw [CoRec (concat res)]
  55
  56 wwExpr :: PlainCoreExpr -> WwM PlainCoreExpr
  57 wwExpr e@(CoVar _) = returnWw e
  58 wwExpr e@(CoLit _) = returnWw e
  59 wwExpr e@(CoCon _ _ _) = returnWw e
  60 wwExpr e@(CoPrim _ _ _) = returnWw e
  61 wwExpr   (CoLam ids e) =
  62         wwExpr e                `thenWw` \ e' ->
  63         returnWw (CoLam ids e')
  64 wwExpr   (CoTyLam tyvar e) =
  65         wwExpr e                `thenWw` \ e' ->
  66         returnWw (CoTyLam tyvar e')
  67 wwExpr   (CoApp f atom) =
  68         wwExpr f                `thenWw` \ f' ->
  69         returnWw (CoApp f atom)
  70 wwExpr   (CoTyApp f ty) =
  71         wwExpr f                `thenWw` \ f' ->
  72         returnWw (CoTyApp f' ty)
  73 wwExpr   (CoSCC lab e) =
  74         wwExpr e                `thenWw` \ e' ->
  75         returnWw (CoSCC lab e')
  76 wwExpr   (CoLet bnds e) =
  77         wwExpr e                `thenWw` \ e' ->
  78         wwBind bnds             `thenWw` \ bnds' ->
  79         returnWw (foldr CoLet e' bnds')
  80 wwExpr   (CoCase e alts) =
  81         wwExpr e                `thenWw` \ e' ->
  82         wwAlts alts             `thenWw` \ alts' ->
  83         returnWw  (CoCase e' alts')
  84
  85 wwAlts (CoAlgAlts alts deflt) =
  86         mapWw (\(con,binders,e) ->
  87                         wwExpr e        `thenWw` \ e' ->
  88                         returnWw (con,binders,e')) alts `thenWw` \ alts' ->
  89         wwDef deflt                                     `thenWw` \ deflt' ->
  90         returnWw (CoAlgAlts alts' deflt)
  91 wwAlts (CoPrimAlts alts deflt) =
  92         mapWw (\(lit,e) ->
  93                         wwExpr e        `thenWw` \ e' ->
  94                         returnWw (lit,e')) alts         `thenWw` \ alts' ->
  95         wwDef deflt                                     `thenWw` \ deflt' ->
  96         returnWw (CoPrimAlts alts' deflt)
  97
  98 wwDef e@CoNoDefault = returnWw e
  99 wwDef  (CoBindDefault v e) =
 100         wwExpr e                                        `thenWw` \ e' ->
 101         returnWw (CoBindDefault v e')
 102 \end{code}
 103
 104 \begin{code}
 105 try_split_bind :: Id -> PlainCoreExpr -> WwM [(Id,PlainCoreExpr)]
 106 try_split_bind id expr =
 107   wwExpr expr                   `thenWw` \ expr' ->
 108   case getFBType (getIdFBTypeInfo id) of
 109     Just (FBType consum prod)
 110         |  FBGoodProd == prod ->
 111 {-      || any (== FBGoodConsum) consum -}
 112       let
 113         (big_args,args,body) = digForLambdas expr'
 114       in
 115         if length args /= length consum   -- funny number of arguments
 116         then returnWw [(id,expr')]
 117         else
 118         -- f /\ t1 .. tn \ v1 .. vn -> e
 119         --      ===>
 120         -- f_wrk /\ t1 .. tn t_new \ v1 .. vn c n -> foldr <exprTy> <nTy> c n e
 121         -- f /\ t1 .. tn \ v1 .. vn
 122         --      -> build exprTy (\ c n -> f_wrk t1 .. tn t_new v1 .. vn c n)
 123         pprTrace "WW:" (ppr PprDebug id) (returnWw ())
 124                                 `thenWw` \ () ->
 125         getUniqueWw             `thenWw` \ ty_new_uq ->
 126         getUniqueWw             `thenWw` \ worker_new_uq ->
 127         getUniqueWw             `thenWw` \ c_new_uq ->
 128         getUniqueWw             `thenWw` \ n_new_uq ->
 129       let
 130         -- The *new* type
 131         n_ty = alpha_ty
 132         n_ty_templ = alpha
 133
 134         (templ,arg_tys,res) = splitTypeWithDictsAsArgs (getIdUniType id)
 135         expr_ty = getListTy res
 136         getListTy res = case res of
 137                          UniData lty [ty] | lty `eqTyCon` listTyCon -> ty
 138                          _ -> panic "Trying to split a non List datatype into Worker/Wrapper"
 139
 140         c_ty       = expr_ty `mkFunTy` (n_ty `mkFunTy` n_ty)
 141         c_ty_templ = expr_ty `mkFunTy` (n_ty_templ `mkFunTy` n_ty_templ)
 142
 143         worker_ty = mkForallTy (templ  ++ [alpha_tv])
 144                         (foldr mkFunTy n_ty_templ (arg_tys++[c_ty_templ,n_ty_templ]))
 145         wrapper_id  = id `replaceIdInfo`
 146                               (getIdInfo id     `addInfo_UF`
 147                                iWantToBeINLINEd UnfoldAlways)
 148         worker_id  = mkWorkerId worker_new_uq id worker_ty
 149                                 noIdInfo
 150                 -- TODO : CHECK if mkWorkerId is thr
 151                 -- right function to use ..
 152         -- Now the bodies
 153
 154         c_id = mkSysLocal SLIT("_fbww") c_new_uq c_ty mkUnknownSrcLoc
 155         n_id = mkSysLocal SLIT("_fbww") n_new_uq n_ty mkUnknownSrcLoc
 156         worker_rhs = foldr CoTyLam
 157                         (mkCoLam (args++[c_id,n_id]) worker_body)
 158                         (big_args ++ [alpha_tyvar])
 159         worker_body = runBuiltinUs (
 160                          mkCoApps (mkCoTyApps (CoVar foldrId) [expr_ty, n_ty])
 161                                   [CoVar c_id,CoVar n_id,body])
 162         wrapper_rhs = foldr CoTyLam
 163                         (mkCoLam (args) wrapper_body)
 164                         big_args
 165         wrapper_body = runBuiltinUs (
 166                  mkCoApps (mkCoTyApp (CoVar buildId) expr_ty)
 167                                 [CoTyLam alpha_tyvar (mkCoLam [c_id,n_id]
 168                 (foldl CoApp
 169                         (mkCoTyApps (CoVar worker_id)
 170                                 [mkTyVarTy t | t <- big_args ++ [alpha_tyvar]])
 171                         (map CoVarAtom (args++[c_id,n_id]))))])
 172
 173       in
 174         if length args /= length arg_tys ||
 175            length big_args /= length templ
 176         then panic "LEN PROBLEM"
 177         else
 178         returnWw  [(worker_id,worker_rhs),(wrapper_id,wrapper_rhs)]
 179     _ -> returnWw [(id,expr')]
 180 \end{code}
 181