2 % (c) The University of Glasgow 2006
3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
6 Desugaring list comprehensions and array comprehensions
9 module DsListComp ( dsListComp, dsPArrComp ) where
11 #include "HsVersions.h"
13 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLocalBinds )
20 import DsMonad -- the monadery used in the desugarer
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).
40 There will be at least one ``qualifier'' in the input.
43 dsListComp :: [LStmt Id]
45 -> Type -- Type of list elements
47 dsListComp lquals body elt_ty
48 = getDOptsDs `thenDs` \dflags ->
50 quals = map unLoc lquals
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)
61 else -- Foldr/build should be enabled, so desugar
62 -- into foldrs and builds
63 newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
65 n_ty = mkTyVarTy n_tyvar
66 c_ty = mkFunTys [elt_ty, n_ty] n_ty
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)
74 where isParallelComp (ParStmt bndrstmtss : _) = True
75 isParallelComp _ = False
78 %************************************************************************
80 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
82 %************************************************************************
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.
93 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
96 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
99 TQ << [ e | b , qs ] ++ L >> =
100 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
103 TQ << [ e | p <- L1, qs ] ++ L2 >> =
109 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
114 "h", "u1", "u2", and "u3" are new variables.
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.
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]
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
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.
143 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
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) ->
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))
154 bndrs_s = map snd stmtss_w_bndrs
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
160 -- Types of (x1,..,xn), (y1,..,yn) etc
161 qual_tys = map mk_bndrs_tys bndrs_s
163 do_list_comp (stmts, bndrs)
164 = dsListComp stmts (mk_hs_tuple_expr bndrs)
167 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
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)
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)
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
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
192 deBindComp pat core_list1 quals body core_list2
194 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
196 -- u1_ty is a [alpha] type, and u2_ty = alpha
197 u2_ty = hsLPatType pat
199 res_ty = exprType core_list2
200 h_ty = u1_ty `mkFunTy` res_ty
202 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
204 -- the "fail" value ...
206 core_fail = App (Var h) (Var u3)
207 letrec_body = App (Var h) core_list1
209 deListComp quals body core_fail `thenDs` \ rest_expr ->
210 matchSimply (Var u2) (StmtCtxt ListComp) pat
211 rest_expr core_fail `thenDs` \ core_match ->
214 Case (Var u1) u1 res_ty
215 [(DataAlt nilDataCon, [], core_list2),
216 (DataAlt consDataCon, [u2, u3], core_match)]
217 -- Increasing order of tag
219 returnDs (Let (Rec [(h, rhs)]) letrec_body)
224 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
225 -- mkZipBind [t1, t2]
226 -- = (zip, \as1:[t1] as2:[t2]
229 -- (a1:as'1) -> case as2 of
231 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
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 ->
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)
244 returnDs (zip_fn, mkLams ass zip_body)
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
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
262 mk_hs_tuple_pat :: [Id] -> LPat Id
263 mk_hs_tuple_pat bs = mkTuplePat (map nlVarPat bs)
267 %************************************************************************
269 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
271 %************************************************************************
273 @dfListComp@ are the rules used with foldr/build turned on:
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
287 dfListComp :: Id -> Id -- 'c' and 'n'
288 -> [Stmt Id] -- the rest of the qual's
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])
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))
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
308 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
309 -- evaluate the two lists
310 = dsLExpr list1 `thenDs` \ core_list1 ->
312 -- find the required type
313 let x_ty = hsLPatType pat
317 -- create some new local id's
318 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
320 -- build rest of the comprehesion
321 dfListComp c_id b quals body `thenDs` \ core_rest ->
323 -- build the pattern match
324 matchSimply (Var x) (StmtCtxt ListComp)
325 pat core_rest (Var b) `thenDs` \ core_expr ->
327 -- now build the outermost foldr, and return
328 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
330 Var foldr_id `App` Type x_ty
332 `App` mkLams [x, b] core_expr
338 %************************************************************************
340 \subsection[DsPArrComp]{Desugaring of array comprehensions}
342 %************************************************************************
346 -- entry point for desugaring a parallel array comprehension
348 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
350 dsPArrComp :: [Stmt Id]
352 -> Type -- Don't use; called with `undefined' below
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,
362 dePArrComp qs body (mkTuplePat []) unitArray
368 dePArrComp :: [Stmt Id]
370 -> LPat Id -- the current generator pattern
371 -> CoreExpr -- the current generator expression
374 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
376 dePArrComp [] e' pa cea =
377 dsLookupGlobalId mapPName `thenDs` \mapP ->
378 let ty = parrElemType cea
380 deLambda ty pa e' `thenDs` \(clam,
382 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
384 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
386 dePArrComp (ExprStmt b _ _ : qs) body pa cea =
387 dsLookupGlobalId filterPName `thenDs` \filterP ->
388 let ty = parrElemType cea
390 deLambda ty pa b `thenDs` \(clam,_) ->
391 dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
393 -- <<[:e' | p <- e, qs:]>> pa ea =
394 -- let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
396 -- <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
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
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]
411 mkLambda ety'cea pa cef `thenDs` \(clam,
413 let ety'cef = ety'ce -- filter doesn't change the element type
414 pa' = mkTuplePat [pa, p]
416 dePArrComp qs body pa' (mkApps (Var crossMapP)
417 [Type ety'cea, Type ety'cef, cea, clam])
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)
423 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
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
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!"
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
443 dePArrComp qs body pa' (mkApps (Var mapP)
444 [Type ty'cea, Type errTy, proj, cea])
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.
450 dePArrComp (ParStmt _ : _) _ _ _ =
451 panic "DsListComp.dePArrComp: malformed comprehension AST"
453 -- <<[:e' | qs | qss:]>> pa ea =
454 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
455 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
457 -- {x_1, ..., x_n} = DV (qs)
459 dePArrParComp qss body =
460 deParStmt qss `thenDs` \(pQss,
462 dePArrComp [] body pQss ceQss
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)
470 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
471 parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
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)
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]
484 parStmts qss pa' cea'
486 -- generate Core corresponding to `\p -> e'
488 deLambda :: Type -- type of the argument
489 -> LPat Id -- argument pattern
490 -> LHsExpr Id -- body
491 -> DsM (CoreExpr, Type)
493 dsLExpr e `thenDs` \ce ->
496 -- generate Core for a lambda pattern match, where the body is already in Core
498 mkLambda :: Type -- type of the argument
499 -> LPat Id -- argument pattern
500 -> CoreExpr -- desugared body
501 -> DsM (CoreExpr, Type)
503 newSysLocalDs ty `thenDs` \v ->
504 let errMsg = "DsListComp.deLambda: internal error!"
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)
511 -- obtain the element type of the parallel array produced by the given Core
514 parrElemType :: CoreExpr -> Type
516 case splitTyConApp_maybe (exprType e) of
517 Just (tycon, [ty]) | tycon == parrTyCon -> ty
519 "DsListComp.parrElemType: not a parallel array type"
521 -- Smart constructor for source tuple patterns
523 mkTuplePat :: [LPat Id] -> LPat Id
524 mkTuplePat [lpat] = lpat
525 mkTuplePat lpats = noLoc $ mkVanillaTuplePat lpats Boxed
527 -- Smart constructor for source tuple expressions
529 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
530 mkExplicitTuple [lexp] = lexp
531 mkExplicitTuple lexps = noLoc $ ExplicitTuple lexps Boxed