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