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
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).
41 There will be at least one ``qualifier'' in the input.
44 dsListComp :: [LStmt Id]
46 -> Type -- Type of list elements
48 dsListComp lquals body elt_ty
49 = getDOptsDs `thenDs` \dflags ->
51 quals = map unLoc lquals
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)
62 else -- Foldr/build should be enabled, so desugar
63 -- into foldrs and builds
64 newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
66 n_ty = mkTyVarTy n_tyvar
67 c_ty = mkFunTys [elt_ty, n_ty] n_ty
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)
75 where isParallelComp (ParStmt bndrstmtss : _) = True
76 isParallelComp _ = False
79 %************************************************************************
81 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
83 %************************************************************************
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.
94 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
97 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
100 TQ << [ e | b , qs ] ++ L >> =
101 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
104 TQ << [ e | p <- L1, qs ] ++ L2 >> =
110 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
115 "h", "u1", "u2", and "u3" are new variables.
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.
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]
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
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.
144 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
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) ->
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))
155 bndrs_s = map snd stmtss_w_bndrs
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
161 -- Types of (x1,..,xn), (y1,..,yn) etc
162 qual_tys = map mk_bndrs_tys bndrs_s
164 do_list_comp (stmts, bndrs)
165 = dsListComp stmts (mk_hs_tuple_expr bndrs)
168 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
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)
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)
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
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
193 deBindComp pat core_list1 quals body core_list2
195 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
197 -- u1_ty is a [alpha] type, and u2_ty = alpha
198 u2_ty = hsLPatType pat
200 res_ty = exprType core_list2
201 h_ty = u1_ty `mkFunTy` res_ty
203 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
205 -- the "fail" value ...
207 core_fail = App (Var h) (Var u3)
208 letrec_body = App (Var h) core_list1
210 deListComp quals body core_fail `thenDs` \ rest_expr ->
211 matchSimply (Var u2) (StmtCtxt ListComp) pat
212 rest_expr core_fail `thenDs` \ core_match ->
215 Case (Var u1) u1 res_ty
216 [(DataAlt nilDataCon, [], core_list2),
217 (DataAlt consDataCon, [u2, u3], core_match)]
218 -- Increasing order of tag
220 returnDs (Let (Rec [(h, rhs)]) letrec_body)
225 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
226 -- mkZipBind [t1, t2]
227 -- = (zip, \as1:[t1] as2:[t2]
230 -- (a1:as'1) -> case as2 of
232 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
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 ->
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)
245 returnDs (zip_fn, mkLams ass zip_body)
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
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
263 mk_hs_tuple_pat :: [Id] -> LPat Id
264 mk_hs_tuple_pat bs = mkTuplePat (map nlVarPat bs)
268 %************************************************************************
270 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
272 %************************************************************************
274 @dfListComp@ are the rules used with foldr/build turned on:
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
288 dfListComp :: Id -> Id -- 'c' and 'n'
289 -> [Stmt Id] -- the rest of the qual's
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])
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))
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
309 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
310 -- evaluate the two lists
311 = dsLExpr list1 `thenDs` \ core_list1 ->
313 -- find the required type
314 let x_ty = hsLPatType pat
318 -- create some new local id's
319 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
321 -- build rest of the comprehesion
322 dfListComp c_id b quals body `thenDs` \ core_rest ->
324 -- build the pattern match
325 matchSimply (Var x) (StmtCtxt ListComp)
326 pat core_rest (Var b) `thenDs` \ core_expr ->
328 -- now build the outermost foldr, and return
329 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
331 Var foldr_id `App` Type x_ty
333 `App` mkLams [x, b] core_expr
339 %************************************************************************
341 \subsection[DsPArrComp]{Desugaring of array comprehensions}
343 %************************************************************************
347 -- entry point for desugaring a parallel array comprehension
349 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
351 dsPArrComp :: [Stmt Id]
353 -> Type -- Don't use; called with `undefined' below
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,
363 dePArrComp qs body (mkTuplePat []) unitArray
369 dePArrComp :: [Stmt Id]
371 -> LPat Id -- the current generator pattern
372 -> CoreExpr -- the current generator expression
375 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
377 dePArrComp [] e' pa cea =
378 dsLookupGlobalId mapPName `thenDs` \mapP ->
379 let ty = parrElemType cea
381 deLambda ty pa e' `thenDs` \(clam,
383 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
385 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
387 dePArrComp (ExprStmt b _ _ : qs) body pa cea =
388 dsLookupGlobalId filterPName `thenDs` \filterP ->
389 let ty = parrElemType cea
391 deLambda ty pa b `thenDs` \(clam,_) ->
392 dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
394 -- <<[:e' | p <- e, qs:]>> pa ea =
395 -- let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
397 -- <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
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
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]
412 mkLambda ety'cea pa cef `thenDs` \(clam,
414 let ety'cef = ety'ce -- filter doesn't change the element type
415 pa' = mkTuplePat [pa, p]
417 dePArrComp qs body pa' (mkApps (Var crossMapP)
418 [Type ety'cea, Type ety'cef, cea, clam])
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)
424 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
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
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!"
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
444 dePArrComp qs body pa' (mkApps (Var mapP)
445 [Type ty'cea, Type errTy, proj, cea])
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.
451 dePArrComp (ParStmt _ : _) _ _ _ =
452 panic "DsListComp.dePArrComp: malformed comprehension AST"
454 -- <<[:e' | qs | qss:]>> pa ea =
455 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
456 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
458 -- {x_1, ..., x_n} = DV (qs)
460 dePArrParComp qss body =
461 deParStmt qss `thenDs` \(pQss,
463 dePArrComp [] body pQss ceQss
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)
471 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
472 parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
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)
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]
485 parStmts qss pa' cea'
487 -- generate Core corresponding to `\p -> e'
489 deLambda :: Type -- type of the argument
490 -> LPat Id -- argument pattern
491 -> LHsExpr Id -- body
492 -> DsM (CoreExpr, Type)
494 dsLExpr e `thenDs` \ce ->
497 -- generate Core for a lambda pattern match, where the body is already in Core
499 mkLambda :: Type -- type of the argument
500 -> LPat Id -- argument pattern
501 -> CoreExpr -- desugared body
502 -> DsM (CoreExpr, Type)
504 newSysLocalDs ty `thenDs` \v ->
505 let errMsg = "DsListComp.deLambda: internal error!"
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)
512 -- obtain the element type of the parallel array produced by the given Core
515 parrElemType :: CoreExpr -> Type
517 case splitTyConApp_maybe (exprType e) of
518 Just (tycon, [ty]) | tycon == parrTyCon -> ty
520 "DsListComp.parrElemType: not a parallel array type"
522 -- Smart constructor for source tuple patterns
524 mkTuplePat :: [LPat Id] -> LPat Id
525 mkTuplePat [lpat] = lpat
526 mkTuplePat lpats = noLoc $ mkVanillaTuplePat lpats Boxed
528 -- Smart constructor for source tuple expressions
530 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
531 mkExplicitTuple [lexp] = lexp
532 mkExplicitTuple lexps = noLoc $ ExplicitTuple lexps Boxed