2 % (c) The University of Glasgow 2006
3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
6 Desugaring list comprehensions and array comprehensions
10 -- The above warning supression flag is a temporary kludge.
11 -- While working on this module you are encouraged to remove it and fix
12 -- any warnings in the module. See
13 -- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
16 module DsListComp ( dsListComp, dsPArrComp ) where
18 #include "HsVersions.h"
20 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLocalBinds )
27 import DsMonad -- the monadery used in the desugarer
43 List comprehensions may be desugared in one of two ways: ``ordinary''
44 (as you would expect if you read SLPJ's book) and ``with foldr/build
45 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
47 There will be at least one ``qualifier'' in the input.
50 dsListComp :: [LStmt Id]
52 -> Type -- Type of list elements
54 dsListComp lquals body elt_ty
55 = getDOptsDs `thenDs` \dflags ->
57 quals = map unLoc lquals
59 if not (dopt Opt_RewriteRules dflags) || dopt Opt_IgnoreInterfacePragmas dflags
60 -- Either rules are switched off, or we are ignoring what there are;
61 -- Either way foldr/build won't happen, so use the more efficient
62 -- Wadler-style desugaring
63 || isParallelComp quals
64 -- Foldr-style desugaring can't handle
65 -- parallel list comprehensions
66 then deListComp quals body (mkNilExpr elt_ty)
68 else -- Foldr/build should be enabled, so desugar
69 -- into foldrs and builds
70 newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
72 n_ty = mkTyVarTy n_tyvar
73 c_ty = mkFunTys [elt_ty, n_ty] n_ty
75 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
76 dfListComp c n quals body `thenDs` \ result ->
77 dsLookupGlobalId buildName `thenDs` \ build_id ->
78 returnDs (Var build_id `App` Type elt_ty
79 `App` mkLams [n_tyvar, c, n] result)
81 where isParallelComp (ParStmt bndrstmtss : _) = True
82 isParallelComp _ = False
85 %************************************************************************
87 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
89 %************************************************************************
91 Just as in Phil's chapter~7 in SLPJ, using the rules for
92 optimally-compiled list comprehensions. This is what Kevin followed
93 as well, and I quite happily do the same. The TQ translation scheme
94 transforms a list of qualifiers (either boolean expressions or
95 generators) into a single expression which implements the list
96 comprehension. Because we are generating 2nd-order polymorphic
97 lambda-calculus, calls to NIL and CONS must be applied to a type
98 argument, as well as their usual value arguments.
100 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
103 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
106 TQ << [ e | b , qs ] ++ L >> =
107 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
110 TQ << [ e | p <- L1, qs ] ++ L2 >> =
116 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
121 "h", "u1", "u2", and "u3" are new variables.
124 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
125 is the TE translation scheme. Note that we carry around the @L@ list
126 already desugared. @dsListComp@ does the top TE rule mentioned above.
128 To the above, we add an additional rule to deal with parallel list
129 comprehensions. The translation goes roughly as follows:
130 [ e | p1 <- e11, let v1 = e12, p2 <- e13
131 | q1 <- e21, let v2 = e22, q2 <- e23]
133 [ e | ((x1, .., xn), (y1, ..., ym)) <-
134 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
135 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
136 where (x1, .., xn) are the variables bound in p1, v1, p2
137 (y1, .., ym) are the variables bound in q1, v2, q2
139 In the translation below, the ParStmt branch translates each parallel branch
140 into a sub-comprehension, and desugars each independently. The resulting lists
141 are fed to a zip function, we create a binding for all the variables bound in all
142 the comprehensions, and then we hand things off the the desugarer for bindings.
143 The zip function is generated here a) because it's small, and b) because then we
144 don't have to deal with arbitrary limits on the number of zip functions in the
145 prelude, nor which library the zip function came from.
146 The introduced tuples are Boxed, but only because I couldn't get it to work
147 with the Unboxed variety.
150 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
152 deListComp (ParStmt stmtss_w_bndrs : quals) body list
153 = mappM do_list_comp stmtss_w_bndrs `thenDs` \ exps ->
154 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
156 -- Deal with [e | pat <- zip l1 .. ln] in example above
157 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
161 bndrs_s = map snd stmtss_w_bndrs
163 -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
164 pat = mkTuplePat pats
165 pats = map mk_hs_tuple_pat bndrs_s
167 -- Types of (x1,..,xn), (y1,..,yn) etc
168 qual_tys = map mk_bndrs_tys bndrs_s
170 do_list_comp (stmts, bndrs)
171 = dsListComp stmts (mk_hs_tuple_expr bndrs)
174 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
176 -- Last: the one to return
177 deListComp [] body list -- Figure 7.4, SLPJ, p 135, rule C above
178 = dsLExpr body `thenDs` \ core_body ->
179 returnDs (mkConsExpr (exprType core_body) core_body list)
181 -- Non-last: must be a guard
182 deListComp (ExprStmt guard _ _ : quals) body list -- rule B above
183 = dsLExpr guard `thenDs` \ core_guard ->
184 deListComp quals body list `thenDs` \ core_rest ->
185 returnDs (mkIfThenElse core_guard core_rest list)
187 -- [e | let B, qs] = let B in [e | qs]
188 deListComp (LetStmt binds : quals) body list
189 = deListComp quals body list `thenDs` \ core_rest ->
190 dsLocalBinds binds core_rest
192 deListComp (BindStmt pat list1 _ _ : quals) body core_list2 -- rule A' above
193 = dsLExpr list1 `thenDs` \ core_list1 ->
194 deBindComp pat core_list1 quals body core_list2
199 deBindComp pat core_list1 quals body core_list2
201 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
203 -- u1_ty is a [alpha] type, and u2_ty = alpha
204 u2_ty = hsLPatType pat
206 res_ty = exprType core_list2
207 h_ty = u1_ty `mkFunTy` res_ty
209 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
211 -- the "fail" value ...
213 core_fail = App (Var h) (Var u3)
214 letrec_body = App (Var h) core_list1
216 deListComp quals body core_fail `thenDs` \ rest_expr ->
217 matchSimply (Var u2) (StmtCtxt ListComp) pat
218 rest_expr core_fail `thenDs` \ core_match ->
221 Case (Var u1) u1 res_ty
222 [(DataAlt nilDataCon, [], core_list2),
223 (DataAlt consDataCon, [u2, u3], core_match)]
224 -- Increasing order of tag
226 returnDs (Let (Rec [(h, rhs)]) letrec_body)
231 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
232 -- mkZipBind [t1, t2]
233 -- = (zip, \as1:[t1] as2:[t2]
236 -- (a1:as'1) -> case as2 of
238 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
241 = mappM newSysLocalDs list_tys `thenDs` \ ass ->
242 mappM newSysLocalDs elt_tys `thenDs` \ as' ->
243 mappM newSysLocalDs list_tys `thenDs` \ as's ->
244 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
246 inner_rhs = mkConsExpr ret_elt_ty
247 (mkCoreTup (map Var as'))
248 (mkVarApps (Var zip_fn) as's)
249 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
251 returnDs (zip_fn, mkLams ass zip_body)
253 list_tys = map mkListTy elt_tys
254 ret_elt_ty = mkCoreTupTy elt_tys
255 list_ret_ty = mkListTy ret_elt_ty
256 zip_fn_ty = mkFunTys list_tys list_ret_ty
258 mk_case (as, a', as') rest
259 = Case (Var as) as list_ret_ty
260 [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
261 (DataAlt consDataCon, [a', as'], rest)]
262 -- Increasing order of tag
263 -- Helper functions that makes an HsTuple only for non-1-sized tuples
264 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
265 mk_hs_tuple_expr [] = nlHsVar unitDataConId
266 mk_hs_tuple_expr [id] = nlHsVar id
267 mk_hs_tuple_expr ids = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
269 mk_hs_tuple_pat :: [Id] -> LPat Id
270 mk_hs_tuple_pat bs = mkTuplePat (map nlVarPat bs)
274 %************************************************************************
276 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
278 %************************************************************************
280 @dfListComp@ are the rules used with foldr/build turned on:
283 TE[ e | ] c n = c e n
284 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
285 TE[ e | p <- l , q ] c n = let
286 f = \ x b -> case x of
294 dfListComp :: Id -> Id -- 'c' and 'n'
295 -> [Stmt Id] -- the rest of the qual's
299 -- Last: the one to return
300 dfListComp c_id n_id [] body
301 = dsLExpr body `thenDs` \ core_body ->
302 returnDs (mkApps (Var c_id) [core_body, Var n_id])
304 -- Non-last: must be a guard
305 dfListComp c_id n_id (ExprStmt guard _ _ : quals) body
306 = dsLExpr guard `thenDs` \ core_guard ->
307 dfListComp c_id n_id quals body `thenDs` \ core_rest ->
308 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
310 dfListComp c_id n_id (LetStmt binds : quals) body
311 -- new in 1.3, local bindings
312 = dfListComp c_id n_id quals body `thenDs` \ core_rest ->
313 dsLocalBinds binds core_rest
315 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
316 -- evaluate the two lists
317 = dsLExpr list1 `thenDs` \ core_list1 ->
319 -- find the required type
320 let x_ty = hsLPatType pat
324 -- create some new local id's
325 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
327 -- build rest of the comprehesion
328 dfListComp c_id b quals body `thenDs` \ core_rest ->
330 -- build the pattern match
331 matchSimply (Var x) (StmtCtxt ListComp)
332 pat core_rest (Var b) `thenDs` \ core_expr ->
334 -- now build the outermost foldr, and return
335 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
337 Var foldr_id `App` Type x_ty
339 `App` mkLams [x, b] core_expr
345 %************************************************************************
347 \subsection[DsPArrComp]{Desugaring of array comprehensions}
349 %************************************************************************
353 -- entry point for desugaring a parallel array comprehension
355 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
357 dsPArrComp :: [Stmt Id]
359 -> Type -- Don't use; called with `undefined' below
361 dsPArrComp [ParStmt qss] body _ = -- parallel comprehension
362 dePArrParComp qss body
363 dsPArrComp qs body _ = -- no ParStmt in `qs'
364 dsLookupGlobalId replicatePName `thenDs` \repP ->
365 let unitArray = mkApps (Var repP) [Type unitTy,
369 dePArrComp qs body (mkTuplePat []) unitArray
375 dePArrComp :: [Stmt Id]
377 -> LPat Id -- the current generator pattern
378 -> CoreExpr -- the current generator expression
381 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
383 dePArrComp [] e' pa cea =
384 dsLookupGlobalId mapPName `thenDs` \mapP ->
385 let ty = parrElemType cea
387 deLambda ty pa e' `thenDs` \(clam,
389 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
391 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
393 dePArrComp (ExprStmt b _ _ : qs) body pa cea =
394 dsLookupGlobalId filterPName `thenDs` \filterP ->
395 let ty = parrElemType cea
397 deLambda ty pa b `thenDs` \(clam,_) ->
398 dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
400 -- <<[:e' | p <- e, qs:]>> pa ea =
401 -- let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
403 -- <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
405 dePArrComp (BindStmt p e _ _ : qs) body pa cea =
406 dsLookupGlobalId filterPName `thenDs` \filterP ->
407 dsLookupGlobalId crossMapPName `thenDs` \crossMapP ->
408 dsLExpr e `thenDs` \ce ->
409 let ety'cea = parrElemType cea
410 ety'ce = parrElemType ce
411 false = Var falseDataConId
412 true = Var trueDataConId
414 newSysLocalDs ety'ce `thenDs` \v ->
415 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
416 let cef = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
418 mkLambda ety'cea pa cef `thenDs` \(clam,
420 let ety'cef = ety'ce -- filter doesn't change the element type
421 pa' = mkTuplePat [pa, p]
423 dePArrComp qs body pa' (mkApps (Var crossMapP)
424 [Type ety'cea, Type ety'cef, cea, clam])
426 -- <<[:e' | let ds, qs:]>> pa ea =
427 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
428 -- (mapP (\v@pa -> let ds in (v, (x_1, ..., x_n))) ea)
430 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
432 dePArrComp (LetStmt ds : qs) body pa cea =
433 dsLookupGlobalId mapPName `thenDs` \mapP ->
434 let xs = map unLoc (collectLocalBinders ds)
435 ty'cea = parrElemType cea
437 newSysLocalDs ty'cea `thenDs` \v ->
438 dsLocalBinds ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
439 newSysLocalDs (exprType clet) `thenDs` \let'v ->
440 let projBody = mkDsLet (NonRec let'v clet) $
441 mkCoreTup [Var v, Var let'v]
442 errTy = exprType projBody
443 errMsg = "DsListComp.dePArrComp: internal error!"
445 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
446 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase ->
447 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
448 proj = mkLams [v] ccase
450 dePArrComp qs body pa' (mkApps (Var mapP)
451 [Type ty'cea, Type errTy, proj, cea])
453 -- The parser guarantees that parallel comprehensions can only appear as
454 -- singeltons qualifier lists, which we already special case in the caller.
455 -- So, encountering one here is a bug.
457 dePArrComp (ParStmt _ : _) _ _ _ =
458 panic "DsListComp.dePArrComp: malformed comprehension AST"
460 -- <<[:e' | qs | qss:]>> pa ea =
461 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
462 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
464 -- {x_1, ..., x_n} = DV (qs)
466 dePArrParComp qss body =
467 deParStmt qss `thenDs` \(pQss,
469 dePArrComp [] body pQss ceQss
472 -- empty parallel statement lists have no source representation
473 panic "DsListComp.dePArrComp: Empty parallel list comprehension"
474 deParStmt ((qs, xs):qss) = -- first statement
475 let res_expr = mkExplicitTuple (map nlHsVar xs)
477 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
478 parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
480 parStmts [] pa cea = return (pa, cea)
481 parStmts ((qs, xs):qss) pa cea = -- subsequent statements (zip'ed)
482 dsLookupGlobalId zipPName `thenDs` \zipP ->
483 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
484 ty'cea = parrElemType cea
485 res_expr = mkExplicitTuple (map nlHsVar xs)
487 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
488 let ty'cqs = parrElemType cqs
489 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
491 parStmts qss pa' cea'
493 -- generate Core corresponding to `\p -> e'
495 deLambda :: Type -- type of the argument
496 -> LPat Id -- argument pattern
497 -> LHsExpr Id -- body
498 -> DsM (CoreExpr, Type)
500 dsLExpr e `thenDs` \ce ->
503 -- generate Core for a lambda pattern match, where the body is already in Core
505 mkLambda :: Type -- type of the argument
506 -> LPat Id -- argument pattern
507 -> CoreExpr -- desugared body
508 -> DsM (CoreExpr, Type)
510 newSysLocalDs ty `thenDs` \v ->
511 let errMsg = "DsListComp.deLambda: internal error!"
514 mkErrorAppDs pAT_ERROR_ID ce'ty errMsg `thenDs` \cerr ->
515 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
516 returnDs (mkLams [v] res, ce'ty)
518 -- obtain the element type of the parallel array produced by the given Core
521 parrElemType :: CoreExpr -> Type
523 case splitTyConApp_maybe (exprType e) of
524 Just (tycon, [ty]) | tycon == parrTyCon -> ty
526 "DsListComp.parrElemType: not a parallel array type"
528 -- Smart constructor for source tuple patterns
530 mkTuplePat :: [LPat Id] -> LPat Id
531 mkTuplePat [lpat] = lpat
532 mkTuplePat lpats = noLoc $ mkVanillaTuplePat lpats Boxed
534 -- Smart constructor for source tuple expressions
536 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
537 mkExplicitTuple [lexp] = lexp
538 mkExplicitTuple lexps = noLoc $ ExplicitTuple lexps Boxed