remove extraneous "+1"; column numbers start at zero
[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 StaticFlags
25 import CoreUtils
26 import Var
27 import Type
28 import TysPrim
29 import TysWiredIn
30 import Match
31 import PrelNames
32 import PrelInfo
33 import SrcLoc
34 import Panic
35 \end{code}
36
37 List comprehensions may be desugared in one of two ways: ``ordinary''
38 (as you would expect if you read SLPJ's book) and ``with foldr/build
39 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
40
41 There will be at least one ``qualifier'' in the input.
42
43 \begin{code}
44 dsListComp :: [LStmt Id] 
45            -> LHsExpr Id
46            -> Type              -- Type of list elements
47            -> DsM CoreExpr
48 dsListComp lquals body 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 body (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 body           `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] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
145
146 deListComp (ParStmt stmtss_w_bndrs : quals) body 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 body 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      = mkTuplePat pats
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 (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 [] body list         -- Figure 7.4, SLPJ, p 135, rule C above
172   = dsLExpr body                `thenDs` \ core_body ->
173     returnDs (mkConsExpr (exprType core_body) core_body list)
174
175         -- Non-last: must be a guard
176 deListComp (ExprStmt guard _ _ : quals) body list       -- rule B above
177   = dsLExpr guard               `thenDs` \ core_guard ->
178     deListComp quals body 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) body list
183   = deListComp quals body list  `thenDs` \ core_rest ->
184     dsLocalBinds binds core_rest
185
186 deListComp (BindStmt pat list1 _ _ : quals) body core_list2 -- rule A' above
187   = dsLExpr list1                   `thenDs` \ core_list1 ->
188     deBindComp pat core_list1 quals body core_list2
189 \end{code}
190
191
192 \begin{code}
193 deBindComp pat core_list1 quals body 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 = hsLPatType 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 body 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 bs  = mkTuplePat (map nlVarPat bs)
265 \end{code}
266
267
268 %************************************************************************
269 %*                                                                      *
270 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
271 %*                                                                      *
272 %************************************************************************
273
274 @dfListComp@ are the rules used with foldr/build turned on:
275
276 \begin{verbatim}
277 TE[ e | ]            c n = c e n
278 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
279 TE[ e | p <- l , q ] c n = let 
280                                 f = \ x b -> case x of
281                                                   p -> TE[ e | q ] c b
282                                                   _ -> b
283                            in
284                            foldr f n l
285 \end{verbatim}
286
287 \begin{code}
288 dfListComp :: Id -> Id                  -- 'c' and 'n'
289            -> [Stmt Id]         -- the rest of the qual's
290            -> LHsExpr Id
291            -> DsM CoreExpr
292
293         -- Last: the one to return
294 dfListComp c_id n_id [] body
295   = dsLExpr body                `thenDs` \ core_body ->
296     returnDs (mkApps (Var c_id) [core_body, Var n_id])
297
298         -- Non-last: must be a guard
299 dfListComp c_id n_id (ExprStmt guard _ _  : quals) body
300   = dsLExpr guard                       `thenDs` \ core_guard ->
301     dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
302     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
303
304 dfListComp c_id n_id (LetStmt binds : quals) body
305   -- new in 1.3, local bindings
306   = dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
307     dsLocalBinds binds core_rest
308
309 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
310     -- evaluate the two lists
311   = dsLExpr list1                       `thenDs` \ core_list1 ->
312
313     -- find the required type
314     let x_ty   = hsLPatType 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 body                `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                 -> LHsExpr Id
353                 -> Type             -- Don't use; called with `undefined' below
354                 -> DsM CoreExpr
355 dsPArrComp [ParStmt qss] body _  =  -- parallel comprehension
356   dePArrParComp qss body
357 dsPArrComp qs            body _  =  -- no ParStmt in `qs'
358   dsLookupGlobalId replicatePName                         `thenDs` \repP ->
359   let unitArray = mkApps (Var repP) [Type unitTy, 
360                                      mkIntExpr 1, 
361                                      mkCoreTup []]
362   in
363   dePArrComp qs body (mkTuplePat []) unitArray
364
365
366
367 -- the work horse
368 --
369 dePArrComp :: [Stmt Id] 
370            -> LHsExpr Id
371            -> LPat Id           -- the current generator pattern
372            -> CoreExpr          -- the current generator expression
373            -> DsM CoreExpr
374 --
375 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
376 --
377 dePArrComp [] e' pa cea =
378   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
379   let ty = parrElemType cea
380   in
381   deLambda ty pa e'                                       `thenDs` \(clam, 
382                                                                      ty'e') ->
383   returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
384 --
385 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
386 --
387 dePArrComp (ExprStmt b _ _ : qs) body pa cea =
388   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
389   let ty = parrElemType cea
390   in
391   deLambda ty pa b                                `thenDs` \(clam,_) ->
392   dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
393 --
394 --  <<[:e' | p <- e, qs:]>> pa ea = 
395 --    let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
396 --    in
397 --    <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
398 --
399 dePArrComp (BindStmt p e _ _ : qs) body pa cea =
400   dsLookupGlobalId filterPName                    `thenDs` \filterP    ->
401   dsLookupGlobalId crossMapPName                  `thenDs` \crossMapP  ->
402   dsLExpr e                                       `thenDs` \ce         ->
403   let ety'cea = parrElemType cea
404       ety'ce  = parrElemType ce
405       false   = Var falseDataConId
406       true    = Var trueDataConId
407   in
408   newSysLocalDs ety'ce                                    `thenDs` \v       ->
409   matchSimply (Var v) (StmtCtxt PArrComp) p true false    `thenDs` \pred    ->
410   let cef = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
411   in
412   mkLambda ety'cea pa cef                                 `thenDs` \(clam, 
413                                                                      _    ) ->
414   let ety'cef = ety'ce              -- filter doesn't change the element type
415       pa'     = mkTuplePat [pa, p]
416   in
417   dePArrComp qs body pa' (mkApps (Var crossMapP) 
418                                  [Type ety'cea, Type ety'cef, cea, clam])
419 --
420 --  <<[:e' | let ds, qs:]>> pa ea = 
421 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n)) 
422 --                    (mapP (\v@pa -> let ds in (v, (x_1, ..., x_n))) ea)
423 --  where
424 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
425 --
426 dePArrComp (LetStmt ds : qs) body pa cea =
427   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
428   let xs     = map unLoc (collectLocalBinders ds)
429       ty'cea = parrElemType cea
430   in
431   newSysLocalDs ty'cea                                    `thenDs` \v       ->
432   dsLocalBinds ds (mkCoreTup (map Var xs))                `thenDs` \clet    ->
433   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
434   let projBody = mkDsLet (NonRec let'v clet) $ 
435                  mkCoreTup [Var v, Var let'v]
436       errTy    = exprType projBody
437       errMsg   = "DsListComp.dePArrComp: internal error!"
438   in
439   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
440   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase   ->
441   let pa'    = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
442       proj   = mkLams [v] ccase
443   in
444   dePArrComp qs body pa' (mkApps (Var mapP) 
445                                  [Type ty'cea, Type errTy, proj, cea])
446 --
447 -- The parser guarantees that parallel comprehensions can only appear as
448 -- singeltons qualifier lists, which we already special case in the caller.
449 -- So, encountering one here is a bug.
450 --
451 dePArrComp (ParStmt _ : _) _ _ _ = 
452   panic "DsListComp.dePArrComp: malformed comprehension AST"
453
454 --  <<[:e' | qs | qss:]>> pa ea = 
455 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n)) 
456 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
457 --    where
458 --      {x_1, ..., x_n} = DV (qs)
459 --
460 dePArrParComp qss body = 
461   deParStmt qss                                         `thenDs` \(pQss, 
462                                                                    ceQss) ->
463   dePArrComp [] body pQss ceQss
464   where
465     deParStmt []             =
466       -- empty parallel statement lists have no source representation
467       panic "DsListComp.dePArrComp: Empty parallel list comprehension"
468     deParStmt ((qs, xs):qss) =          -- first statement
469       let res_expr = mkExplicitTuple (map nlHsVar xs)
470       in
471       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
472       parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
473     ---
474     parStmts []             pa cea = return (pa, cea)
475     parStmts ((qs, xs):qss) pa cea =    -- subsequent statements (zip'ed)
476       dsLookupGlobalId zipPName                           `thenDs` \zipP    ->
477       let pa'      = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
478           ty'cea   = parrElemType cea
479           res_expr = mkExplicitTuple (map nlHsVar xs)
480       in
481       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
482       let ty'cqs = parrElemType cqs
483           cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
484       in
485       parStmts qss pa' cea'
486
487 -- generate Core corresponding to `\p -> e'
488 --
489 deLambda :: Type                        -- type of the argument
490           -> LPat Id                    -- argument pattern
491           -> LHsExpr Id                 -- body
492           -> DsM (CoreExpr, Type)
493 deLambda ty p e =
494   dsLExpr e                                               `thenDs` \ce      ->
495   mkLambda ty p ce
496
497 -- generate Core for a lambda pattern match, where the body is already in Core
498 --
499 mkLambda :: Type                        -- type of the argument
500          -> LPat Id                     -- argument pattern
501          -> CoreExpr                    -- desugared body
502          -> DsM (CoreExpr, Type)
503 mkLambda ty p ce =
504   newSysLocalDs ty                                        `thenDs` \v       ->
505   let errMsg = "DsListComp.deLambda: internal error!"
506       ce'ty  = exprType ce
507   in
508   mkErrorAppDs pAT_ERROR_ID ce'ty errMsg                  `thenDs` \cerr    -> 
509   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
510   returnDs (mkLams [v] res, ce'ty)
511
512 -- obtain the element type of the parallel array produced by the given Core
513 -- expression
514 --
515 parrElemType   :: CoreExpr -> Type
516 parrElemType e  = 
517   case splitTyConApp_maybe (exprType e) of
518     Just (tycon, [ty]) | tycon == parrTyCon -> ty
519     _                                                     -> panic
520       "DsListComp.parrElemType: not a parallel array type"
521
522 -- Smart constructor for source tuple patterns
523 --
524 mkTuplePat :: [LPat Id] -> LPat Id
525 mkTuplePat [lpat] = lpat
526 mkTuplePat lpats  = noLoc $ mkVanillaTuplePat lpats Boxed
527
528 -- Smart constructor for source tuple expressions
529 --
530 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
531 mkExplicitTuple [lexp] = lexp
532 mkExplicitTuple lexps  = noLoc $ ExplicitTuple lexps Boxed
533 \end{code}