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