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