7eb62ffa386b8780d1e05c816ba2658f9417f589
[ghc-hetmet.git] / ghc / compiler / deSugar / DsListComp.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 \section[DsListComp]{Desugaring list comprehensions and array comprehensions}
5
6 \begin{code}
7 module DsListComp ( dsListComp, dsPArrComp ) where
8
9 #include "HsVersions.h"
10
11 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLocalBinds )
12
13 import BasicTypes       ( Boxity(..) )
14 import HsSyn
15 import TcHsSyn          ( hsPatType )
16 import CoreSyn
17
18 import DsMonad          -- the monadery used in the desugarer
19 import DsUtils
20
21 import DynFlags         ( DynFlag(..), dopt )
22 import StaticFlags      ( opt_RulesOff )
23 import CoreUtils        ( exprType, mkIfThenElse )
24 import Id               ( idType )
25 import Var              ( Id )
26 import Type             ( mkTyVarTy, mkFunTys, mkFunTy, Type,
27                           splitTyConApp_maybe )
28 import TysPrim          ( alphaTyVar )
29 import TysWiredIn       ( nilDataCon, consDataCon, trueDataConId, falseDataConId, 
30                           unitDataConId, unitTy, mkListTy, parrTyCon )
31 import Match            ( matchSimply )
32 import PrelNames        ( foldrName, buildName, replicatePName, mapPName, 
33                           filterPName, zipPName, crossPName ) 
34 import PrelInfo         ( pAT_ERROR_ID )
35 import SrcLoc           ( noLoc, unLoc )
36 import Panic            ( panic )
37 \end{code}
38
39 List comprehensions may be desugared in one of two ways: ``ordinary''
40 (as you would expect if you read SLPJ's book) and ``with foldr/build
41 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
42
43 There will be at least one ``qualifier'' in the input.
44
45 \begin{code}
46 dsListComp :: [LStmt Id] 
47            -> LHsExpr Id
48            -> Type              -- Type of list elements
49            -> DsM CoreExpr
50 dsListComp lquals body elt_ty
51   = getDOptsDs  `thenDs` \dflags ->
52     let
53         quals = map unLoc lquals
54     in
55     if opt_RulesOff || dopt Opt_IgnoreInterfacePragmas dflags
56         -- Either rules are switched off, or we are ignoring what there are;
57         -- Either way foldr/build won't happen, so use the more efficient
58         -- Wadler-style desugaring
59         || isParallelComp quals
60                 -- Foldr-style desugaring can't handle
61                 -- parallel list comprehensions
62         then deListComp quals body (mkNilExpr elt_ty)
63
64    else         -- Foldr/build should be enabled, so desugar 
65                 -- into foldrs and builds
66     newTyVarsDs [alphaTyVar]    `thenDs` \ [n_tyvar] ->
67     let
68         n_ty = mkTyVarTy n_tyvar
69         c_ty = mkFunTys [elt_ty, n_ty] n_ty
70     in
71     newSysLocalsDs [c_ty,n_ty]          `thenDs` \ [c, n] ->
72     dfListComp c n quals body           `thenDs` \ result ->
73     dsLookupGlobalId buildName  `thenDs` \ build_id ->
74     returnDs (Var build_id `App` Type elt_ty 
75                            `App` mkLams [n_tyvar, c, n] result)
76
77   where isParallelComp (ParStmt bndrstmtss : _) = True
78         isParallelComp _                        = False
79 \end{code}
80
81 %************************************************************************
82 %*                                                                      *
83 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
84 %*                                                                      *
85 %************************************************************************
86
87 Just as in Phil's chapter~7 in SLPJ, using the rules for
88 optimally-compiled list comprehensions.  This is what Kevin followed
89 as well, and I quite happily do the same.  The TQ translation scheme
90 transforms a list of qualifiers (either boolean expressions or
91 generators) into a single expression which implements the list
92 comprehension.  Because we are generating 2nd-order polymorphic
93 lambda-calculus, calls to NIL and CONS must be applied to a type
94 argument, as well as their usual value arguments.
95 \begin{verbatim}
96 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
97
98 (Rule C)
99 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
100
101 (Rule B)
102 TQ << [ e | b , qs ] ++ L >> =
103     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
104
105 (Rule A')
106 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
107   letrec
108     h = \ u1 ->
109           case u1 of
110             []        ->  TE << L2 >>
111             (u2 : u3) ->
112                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
113                     [] (h u3)
114   in
115     h ( TE << L1 >> )
116
117 "h", "u1", "u2", and "u3" are new variables.
118 \end{verbatim}
119
120 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
121 is the TE translation scheme.  Note that we carry around the @L@ list
122 already desugared.  @dsListComp@ does the top TE rule mentioned above.
123
124 To the above, we add an additional rule to deal with parallel list
125 comprehensions.  The translation goes roughly as follows:
126      [ e | p1 <- e11, let v1 = e12, p2 <- e13
127          | q1 <- e21, let v2 = e22, q2 <- e23]
128      =>
129      [ e | ((x1, .., xn), (y1, ..., ym)) <-
130                zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
131                    [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
132 where (x1, .., xn) are the variables bound in p1, v1, p2
133       (y1, .., ym) are the variables bound in q1, v2, q2
134
135 In the translation below, the ParStmt branch translates each parallel branch
136 into a sub-comprehension, and desugars each independently.  The resulting lists
137 are fed to a zip function, we create a binding for all the variables bound in all
138 the comprehensions, and then we hand things off the the desugarer for bindings.
139 The zip function is generated here a) because it's small, and b) because then we
140 don't have to deal with arbitrary limits on the number of zip functions in the
141 prelude, nor which library the zip function came from.
142 The introduced tuples are Boxed, but only because I couldn't get it to work
143 with the Unboxed variety.
144
145 \begin{code}
146 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
147
148 deListComp (ParStmt stmtss_w_bndrs : quals) body list
149   = mappM do_list_comp stmtss_w_bndrs   `thenDs` \ exps ->
150     mkZipBind qual_tys                  `thenDs` \ (zip_fn, zip_rhs) ->
151
152         -- Deal with [e | pat <- zip l1 .. ln] in example above
153     deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps)) 
154                    quals body list
155
156   where 
157         bndrs_s = map snd stmtss_w_bndrs
158
159         -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
160         pat      = noLoc (TuplePat pats Boxed)
161         pats     = map mk_hs_tuple_pat bndrs_s
162
163         -- Types of (x1,..,xn), (y1,..,yn) etc
164         qual_tys = map mk_bndrs_tys bndrs_s
165
166         do_list_comp (stmts, bndrs)
167           = dsListComp stmts (mk_hs_tuple_expr bndrs)
168                        (mk_bndrs_tys bndrs)
169
170         mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
171
172         -- Last: the one to return
173 deListComp [] body list         -- Figure 7.4, SLPJ, p 135, rule C above
174   = dsLExpr body                `thenDs` \ core_body ->
175     returnDs (mkConsExpr (exprType core_body) core_body list)
176
177         -- Non-last: must be a guard
178 deListComp (ExprStmt guard _ _ : quals) body list       -- rule B above
179   = dsLExpr guard               `thenDs` \ core_guard ->
180     deListComp quals body list  `thenDs` \ core_rest ->
181     returnDs (mkIfThenElse core_guard core_rest list)
182
183 -- [e | let B, qs] = let B in [e | qs]
184 deListComp (LetStmt binds : quals) body list
185   = deListComp quals body list  `thenDs` \ core_rest ->
186     dsLocalBinds binds core_rest
187
188 deListComp (BindStmt pat list1 _ _ : quals) body core_list2 -- rule A' above
189   = dsLExpr list1                   `thenDs` \ core_list1 ->
190     deBindComp pat core_list1 quals body core_list2
191 \end{code}
192
193
194 \begin{code}
195 deBindComp pat core_list1 quals body core_list2
196   = let
197         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
198
199         -- u1_ty is a [alpha] type, and u2_ty = alpha
200         u2_ty = hsPatType pat
201
202         res_ty = exprType core_list2
203         h_ty   = u1_ty `mkFunTy` res_ty
204     in
205     newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]  `thenDs` \ [h, u1, u2, u3] ->
206
207     -- the "fail" value ...
208     let
209         core_fail   = App (Var h) (Var u3)
210         letrec_body = App (Var h) core_list1
211     in
212     deListComp quals body core_fail             `thenDs` \ rest_expr ->
213     matchSimply (Var u2) (StmtCtxt ListComp) pat
214                 rest_expr core_fail             `thenDs` \ core_match ->
215     let
216         rhs = Lam u1 $
217               Case (Var u1) u1 res_ty
218                    [(DataAlt nilDataCon,  [],       core_list2),
219                     (DataAlt consDataCon, [u2, u3], core_match)]
220                         -- Increasing order of tag
221     in
222     returnDs (Let (Rec [(h, rhs)]) letrec_body)
223 \end{code}
224
225
226 \begin{code}
227 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
228 -- mkZipBind [t1, t2] 
229 -- = (zip, \as1:[t1] as2:[t2] 
230 --         -> case as1 of 
231 --              [] -> []
232 --              (a1:as'1) -> case as2 of
233 --                              [] -> []
234 --                              (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
235
236 mkZipBind elt_tys 
237   = mappM newSysLocalDs  list_tys       `thenDs` \ ass ->
238     mappM newSysLocalDs  elt_tys        `thenDs` \ as' ->
239     mappM newSysLocalDs  list_tys       `thenDs` \ as's ->
240     newSysLocalDs zip_fn_ty             `thenDs` \ zip_fn ->
241     let 
242         inner_rhs = mkConsExpr ret_elt_ty 
243                         (mkCoreTup (map Var as'))
244                         (mkVarApps (Var zip_fn) as's)
245         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
246     in
247     returnDs (zip_fn, mkLams ass zip_body)
248   where
249     list_tys    = map mkListTy elt_tys
250     ret_elt_ty  = mkCoreTupTy elt_tys
251     list_ret_ty = mkListTy ret_elt_ty
252     zip_fn_ty   = mkFunTys list_tys list_ret_ty
253
254     mk_case (as, a', as') rest
255           = Case (Var as) as list_ret_ty
256                   [(DataAlt nilDataCon,  [],        mkNilExpr ret_elt_ty),
257                    (DataAlt consDataCon, [a', as'], rest)]
258                         -- Increasing order of tag
259 -- Helper functions that makes an HsTuple only for non-1-sized tuples
260 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
261 mk_hs_tuple_expr []   = nlHsVar unitDataConId
262 mk_hs_tuple_expr [id] = nlHsVar id
263 mk_hs_tuple_expr ids  = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
264
265 mk_hs_tuple_pat :: [Id] -> LPat Id
266 mk_hs_tuple_pat [b] = nlVarPat b
267 mk_hs_tuple_pat bs  = noLoc $ TuplePat (map nlVarPat bs) Boxed
268 \end{code}
269
270
271 %************************************************************************
272 %*                                                                      *
273 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
274 %*                                                                      *
275 %************************************************************************
276
277 @dfListComp@ are the rules used with foldr/build turned on:
278
279 \begin{verbatim}
280 TE[ e | ]            c n = c e n
281 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
282 TE[ e | p <- l , q ] c n = let 
283                                 f = \ x b -> case x of
284                                                   p -> TE[ e | q ] c b
285                                                   _ -> b
286                            in
287                            foldr f n l
288 \end{verbatim}
289
290 \begin{code}
291 dfListComp :: Id -> Id                  -- 'c' and 'n'
292            -> [Stmt Id]         -- the rest of the qual's
293            -> LHsExpr Id
294            -> DsM CoreExpr
295
296         -- Last: the one to return
297 dfListComp c_id n_id [] body
298   = dsLExpr body                `thenDs` \ core_body ->
299     returnDs (mkApps (Var c_id) [core_body, Var n_id])
300
301         -- Non-last: must be a guard
302 dfListComp c_id n_id (ExprStmt guard _ _  : quals) body
303   = dsLExpr guard                       `thenDs` \ core_guard ->
304     dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
305     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
306
307 dfListComp c_id n_id (LetStmt binds : quals) body
308   -- new in 1.3, local bindings
309   = dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
310     dsLocalBinds binds core_rest
311
312 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
313     -- evaluate the two lists
314   = dsLExpr list1                       `thenDs` \ core_list1 ->
315
316     -- find the required type
317     let x_ty   = hsPatType pat
318         b_ty   = idType n_id
319     in
320
321     -- create some new local id's
322     newSysLocalsDs [b_ty,x_ty]                  `thenDs` \ [b,x] ->
323
324     -- build rest of the comprehesion
325     dfListComp c_id b quals body                `thenDs` \ core_rest ->
326
327     -- build the pattern match
328     matchSimply (Var x) (StmtCtxt ListComp)
329                 pat core_rest (Var b)           `thenDs` \ core_expr ->
330
331     -- now build the outermost foldr, and return
332     dsLookupGlobalId foldrName          `thenDs` \ foldr_id ->
333     returnDs (
334       Var foldr_id `App` Type x_ty 
335                    `App` Type b_ty
336                    `App` mkLams [x, b] core_expr
337                    `App` Var n_id
338                    `App` core_list1
339     )
340 \end{code}
341
342 %************************************************************************
343 %*                                                                      *
344 \subsection[DsPArrComp]{Desugaring of array comprehensions}
345 %*                                                                      *
346 %************************************************************************
347
348 \begin{code}
349
350 -- entry point for desugaring a parallel array comprehension
351 --
352 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
353 --
354 dsPArrComp      :: [Stmt Id] 
355                 -> LHsExpr Id
356                 -> Type             -- Don't use; called with `undefined' below
357                 -> DsM CoreExpr
358 dsPArrComp qs body _  =
359   dsLookupGlobalId replicatePName                         `thenDs` \repP ->
360   let unitArray = mkApps (Var repP) [Type unitTy, 
361                                      mkIntExpr 1, 
362                                      mkCoreTup []]
363   in
364   dePArrComp qs body (mkTuplePat []) unitArray
365
366 -- the work horse
367 --
368 dePArrComp :: [Stmt Id] 
369            -> LHsExpr Id
370            -> LPat Id           -- the current generator pattern
371            -> CoreExpr          -- the current generator expression
372            -> DsM CoreExpr
373 --
374 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
375 --
376 dePArrComp [] e' pa cea =
377   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
378   let ty = parrElemType cea
379   in
380   deLambda ty pa e'                                       `thenDs` \(clam, 
381                                                                      ty'e') ->
382   returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
383 --
384 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
385 --
386 dePArrComp (ExprStmt b _ _ : qs) body pa cea =
387   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
388   let ty = parrElemType cea
389   in
390   deLambda ty pa b                                `thenDs` \(clam,_) ->
391   dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
392 --
393 --  <<[:e' | p <- e, qs:]>> pa ea = 
394 --    let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
395 --    in
396 --    <<[:e' | qs:]>> (pa, p) (crossP ea ef)
397 --
398 dePArrComp (BindStmt p e _ _ : qs) body pa cea =
399   dsLookupGlobalId filterPName                    `thenDs` \filterP ->
400   dsLookupGlobalId crossPName                     `thenDs` \crossP  ->
401   dsLExpr e                                       `thenDs` \ce      ->
402   let ty'cea = parrElemType cea
403       ty'ce  = parrElemType ce
404       false  = Var falseDataConId
405       true   = Var trueDataConId
406   in
407   newSysLocalDs ty'ce                                     `thenDs` \v       ->
408   matchSimply (Var v) (StmtCtxt PArrComp) p true false    `thenDs` \pred    ->
409   let cef    = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
410       ty'cef = ty'ce                            -- filterP preserves the type
411       pa'    = mkTuplePat [pa, p]
412   in
413   dePArrComp qs body pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
414 --
415 --  <<[:e' | let ds, qs:]>> pa ea = 
416 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n)) 
417 --                    (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
418 --  where
419 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
420 --
421 dePArrComp (LetStmt ds : qs) body pa cea =
422   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
423   let xs     = map unLoc (collectLocalBinders ds)
424       ty'cea = parrElemType cea
425   in
426   newSysLocalDs ty'cea                                    `thenDs` \v       ->
427   dsLocalBinds ds (mkCoreTup (map Var xs))                `thenDs` \clet    ->
428   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
429   let projBody = mkDsLet (NonRec let'v clet) $ 
430                  mkCoreTup [Var v, Var let'v]
431       errTy    = exprType projBody
432       errMsg   = "DsListComp.dePArrComp: internal error!"
433   in
434   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
435   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase   ->
436   let pa'    = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
437       proj   = mkLams [v] ccase
438   in
439   dePArrComp qs body pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
440 --
441 --  <<[:e' | qs | qss:]>> pa ea = 
442 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n)) 
443 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
444 --    where
445 --      {x_1, ..., x_n} = DV (qs)
446 --
447 dePArrComp (ParStmt qss : qs) body pa cea = 
448   dsLookupGlobalId crossPName                           `thenDs` \crossP  ->
449   deParStmt qss                                         `thenDs` \(pQss, 
450                                                                    ceQss) ->
451   let ty'cea   = parrElemType cea
452       ty'ceQss = parrElemType ceQss
453       pa'      = mkTuplePat [pa, pQss]
454   in
455   dePArrComp qs body pa' (mkApps (Var crossP) [Type ty'cea, Type ty'ceQss, 
456                                                cea, ceQss])
457   where
458     deParStmt []             =
459       -- empty parallel statement lists have not source representation
460       panic "DsListComp.dePArrComp: Empty parallel list comprehension"
461     deParStmt ((qs, xs):qss) =          -- first statement
462       let res_expr = mkExplicitTuple (map nlHsVar xs)
463       in
464       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
465       parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
466     ---
467     parStmts []             pa cea = return (pa, cea)
468     parStmts ((qs, xs):qss) pa cea =    -- subsequent statements (zip'ed)
469       dsLookupGlobalId zipPName                           `thenDs` \zipP    ->
470       let pa'      = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
471           ty'cea   = parrElemType cea
472           res_expr = mkExplicitTuple (map nlHsVar xs)
473       in
474       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
475       let ty'cqs = parrElemType cqs
476           cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
477       in
478       parStmts qss pa' cea'
479
480 -- generate Core corresponding to `\p -> e'
481 --
482 deLambda        :: Type                 -- type of the argument
483                 -> LPat Id              -- argument pattern
484                 -> LHsExpr Id           -- body
485                 -> DsM (CoreExpr, Type)
486 deLambda ty p e  =
487   newSysLocalDs ty                                        `thenDs` \v       ->
488   dsLExpr e                                               `thenDs` \ce      ->
489   let errTy    = exprType ce
490       errMsg   = "DsListComp.deLambda: internal error!"
491   in
492   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    -> 
493   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
494   returnDs (mkLams [v] res, errTy)
495
496 -- obtain the element type of the parallel array produced by the given Core
497 -- expression
498 --
499 parrElemType   :: CoreExpr -> Type
500 parrElemType e  = 
501   case splitTyConApp_maybe (exprType e) of
502     Just (tycon, [ty]) | tycon == parrTyCon -> ty
503     _                                                     -> panic
504       "DsListComp.parrElemType: not a parallel array type"
505
506 -- Smart constructor for source tuple patterns
507 --
508 mkTuplePat :: [LPat id] -> LPat id
509 mkTuplePat [lpat] = lpat
510 mkTuplePat lpats  = noLoc $ TuplePat lpats Boxed
511
512 -- Smart constructor for source tuple expressions
513 --
514 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
515 mkExplicitTuple [lexp] = lexp
516 mkExplicitTuple lexps  = noLoc $ ExplicitTuple lexps Boxed
517 \end{code}