[project @ 1996-01-08 20:28:12 by partain]
[ghc.git] / 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