2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[DsListComp]{Desugaring list comprehensions and array comprehensions}
7 module DsListComp ( dsListComp, dsPArrComp ) where
9 #include "HsVersions.h"
11 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLocalBinds )
13 import BasicTypes ( Boxity(..) )
15 import TcHsSyn ( hsPatType )
18 import DsMonad -- the monadery used in the desugarer
21 import DynFlags ( DynFlag(..), dopt )
22 import StaticFlags ( opt_RulesOff )
23 import CoreUtils ( exprType, mkIfThenElse )
26 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type,
28 import TysPrim ( alphaTyVar )
29 import TysWiredIn ( nilDataCon, consDataCon, trueDataConId, falseDataConId,
30 unitDataConId, unitTy, mkListTy, parrTyCon )
31 import Match ( matchSimply )
32 import PrelNames ( foldrName, buildName, replicatePName, mapPName,
33 filterPName, zipPName, crossPName )
34 import PrelInfo ( pAT_ERROR_ID )
35 import SrcLoc ( noLoc, unLoc )
36 import Panic ( panic )
39 List comprehensions may be desugared in one of two ways: ``ordinary''
40 (as you would expect if you read SLPJ's book) and ``with foldr/build
41 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
43 There will be at least one ``qualifier'' in the input.
46 dsListComp :: [LStmt Id]
48 -> Type -- Type of list elements
50 dsListComp lquals body elt_ty
51 = getDOptsDs `thenDs` \dflags ->
53 quals = map unLoc lquals
55 if opt_RulesOff || dopt Opt_IgnoreInterfacePragmas dflags
56 -- Either rules are switched off, or we are ignoring what there are;
57 -- Either way foldr/build won't happen, so use the more efficient
58 -- Wadler-style desugaring
59 || isParallelComp quals
60 -- Foldr-style desugaring can't handle
61 -- parallel list comprehensions
62 then deListComp quals body (mkNilExpr elt_ty)
64 else -- Foldr/build should be enabled, so desugar
65 -- into foldrs and builds
66 newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
68 n_ty = mkTyVarTy n_tyvar
69 c_ty = mkFunTys [elt_ty, n_ty] n_ty
71 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
72 dfListComp c n quals body `thenDs` \ result ->
73 dsLookupGlobalId buildName `thenDs` \ build_id ->
74 returnDs (Var build_id `App` Type elt_ty
75 `App` mkLams [n_tyvar, c, n] result)
77 where isParallelComp (ParStmt bndrstmtss : _) = True
78 isParallelComp _ = False
81 %************************************************************************
83 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
85 %************************************************************************
87 Just as in Phil's chapter~7 in SLPJ, using the rules for
88 optimally-compiled list comprehensions. This is what Kevin followed
89 as well, and I quite happily do the same. The TQ translation scheme
90 transforms a list of qualifiers (either boolean expressions or
91 generators) into a single expression which implements the list
92 comprehension. Because we are generating 2nd-order polymorphic
93 lambda-calculus, calls to NIL and CONS must be applied to a type
94 argument, as well as their usual value arguments.
96 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
99 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
102 TQ << [ e | b , qs ] ++ L >> =
103 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
106 TQ << [ e | p <- L1, qs ] ++ L2 >> =
112 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
117 "h", "u1", "u2", and "u3" are new variables.
120 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
121 is the TE translation scheme. Note that we carry around the @L@ list
122 already desugared. @dsListComp@ does the top TE rule mentioned above.
124 To the above, we add an additional rule to deal with parallel list
125 comprehensions. The translation goes roughly as follows:
126 [ e | p1 <- e11, let v1 = e12, p2 <- e13
127 | q1 <- e21, let v2 = e22, q2 <- e23]
129 [ e | ((x1, .., xn), (y1, ..., ym)) <-
130 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
131 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
132 where (x1, .., xn) are the variables bound in p1, v1, p2
133 (y1, .., ym) are the variables bound in q1, v2, q2
135 In the translation below, the ParStmt branch translates each parallel branch
136 into a sub-comprehension, and desugars each independently. The resulting lists
137 are fed to a zip function, we create a binding for all the variables bound in all
138 the comprehensions, and then we hand things off the the desugarer for bindings.
139 The zip function is generated here a) because it's small, and b) because then we
140 don't have to deal with arbitrary limits on the number of zip functions in the
141 prelude, nor which library the zip function came from.
142 The introduced tuples are Boxed, but only because I couldn't get it to work
143 with the Unboxed variety.
146 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
148 deListComp (ParStmt stmtss_w_bndrs : quals) body list
149 = mappM do_list_comp stmtss_w_bndrs `thenDs` \ exps ->
150 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
152 -- Deal with [e | pat <- zip l1 .. ln] in example above
153 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
157 bndrs_s = map snd stmtss_w_bndrs
159 -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
160 pat = noLoc (TuplePat pats Boxed)
161 pats = map mk_hs_tuple_pat bndrs_s
163 -- Types of (x1,..,xn), (y1,..,yn) etc
164 qual_tys = map mk_bndrs_tys bndrs_s
166 do_list_comp (stmts, bndrs)
167 = dsListComp stmts (mk_hs_tuple_expr bndrs)
170 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
172 -- Last: the one to return
173 deListComp [] body list -- Figure 7.4, SLPJ, p 135, rule C above
174 = dsLExpr body `thenDs` \ core_body ->
175 returnDs (mkConsExpr (exprType core_body) core_body list)
177 -- Non-last: must be a guard
178 deListComp (ExprStmt guard _ _ : quals) body list -- rule B above
179 = dsLExpr guard `thenDs` \ core_guard ->
180 deListComp quals body list `thenDs` \ core_rest ->
181 returnDs (mkIfThenElse core_guard core_rest list)
183 -- [e | let B, qs] = let B in [e | qs]
184 deListComp (LetStmt binds : quals) body list
185 = deListComp quals body list `thenDs` \ core_rest ->
186 dsLocalBinds binds core_rest
188 deListComp (BindStmt pat list1 _ _ : quals) body core_list2 -- rule A' above
189 = dsLExpr list1 `thenDs` \ core_list1 ->
190 deBindComp pat core_list1 quals body core_list2
195 deBindComp pat core_list1 quals body core_list2
197 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
199 -- u1_ty is a [alpha] type, and u2_ty = alpha
200 u2_ty = hsPatType pat
202 res_ty = exprType core_list2
203 h_ty = u1_ty `mkFunTy` res_ty
205 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
207 -- the "fail" value ...
209 core_fail = App (Var h) (Var u3)
210 letrec_body = App (Var h) core_list1
212 deListComp quals body core_fail `thenDs` \ rest_expr ->
213 matchSimply (Var u2) (StmtCtxt ListComp) pat
214 rest_expr core_fail `thenDs` \ core_match ->
217 Case (Var u1) u1 res_ty
218 [(DataAlt nilDataCon, [], core_list2),
219 (DataAlt consDataCon, [u2, u3], core_match)]
220 -- Increasing order of tag
222 returnDs (Let (Rec [(h, rhs)]) letrec_body)
227 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
228 -- mkZipBind [t1, t2]
229 -- = (zip, \as1:[t1] as2:[t2]
232 -- (a1:as'1) -> case as2 of
234 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
237 = mappM newSysLocalDs list_tys `thenDs` \ ass ->
238 mappM newSysLocalDs elt_tys `thenDs` \ as' ->
239 mappM newSysLocalDs list_tys `thenDs` \ as's ->
240 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
242 inner_rhs = mkConsExpr ret_elt_ty
243 (mkCoreTup (map Var as'))
244 (mkVarApps (Var zip_fn) as's)
245 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
247 returnDs (zip_fn, mkLams ass zip_body)
249 list_tys = map mkListTy elt_tys
250 ret_elt_ty = mkCoreTupTy elt_tys
251 list_ret_ty = mkListTy ret_elt_ty
252 zip_fn_ty = mkFunTys list_tys list_ret_ty
254 mk_case (as, a', as') rest
255 = Case (Var as) as list_ret_ty
256 [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
257 (DataAlt consDataCon, [a', as'], rest)]
258 -- Increasing order of tag
259 -- Helper functions that makes an HsTuple only for non-1-sized tuples
260 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
261 mk_hs_tuple_expr [] = nlHsVar unitDataConId
262 mk_hs_tuple_expr [id] = nlHsVar id
263 mk_hs_tuple_expr ids = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
265 mk_hs_tuple_pat :: [Id] -> LPat Id
266 mk_hs_tuple_pat [b] = nlVarPat b
267 mk_hs_tuple_pat bs = noLoc $ TuplePat (map nlVarPat bs) Boxed
271 %************************************************************************
273 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
275 %************************************************************************
277 @dfListComp@ are the rules used with foldr/build turned on:
280 TE[ e | ] c n = c e n
281 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
282 TE[ e | p <- l , q ] c n = let
283 f = \ x b -> case x of
291 dfListComp :: Id -> Id -- 'c' and 'n'
292 -> [Stmt Id] -- the rest of the qual's
296 -- Last: the one to return
297 dfListComp c_id n_id [] body
298 = dsLExpr body `thenDs` \ core_body ->
299 returnDs (mkApps (Var c_id) [core_body, Var n_id])
301 -- Non-last: must be a guard
302 dfListComp c_id n_id (ExprStmt guard _ _ : quals) body
303 = dsLExpr guard `thenDs` \ core_guard ->
304 dfListComp c_id n_id quals body `thenDs` \ core_rest ->
305 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
307 dfListComp c_id n_id (LetStmt binds : quals) body
308 -- new in 1.3, local bindings
309 = dfListComp c_id n_id quals body `thenDs` \ core_rest ->
310 dsLocalBinds binds core_rest
312 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body
313 -- evaluate the two lists
314 = dsLExpr list1 `thenDs` \ core_list1 ->
316 -- find the required type
317 let x_ty = hsPatType pat
321 -- create some new local id's
322 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
324 -- build rest of the comprehesion
325 dfListComp c_id b quals body `thenDs` \ core_rest ->
327 -- build the pattern match
328 matchSimply (Var x) (StmtCtxt ListComp)
329 pat core_rest (Var b) `thenDs` \ core_expr ->
331 -- now build the outermost foldr, and return
332 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
334 Var foldr_id `App` Type x_ty
336 `App` mkLams [x, b] core_expr
342 %************************************************************************
344 \subsection[DsPArrComp]{Desugaring of array comprehensions}
346 %************************************************************************
350 -- entry point for desugaring a parallel array comprehension
352 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
354 dsPArrComp :: [Stmt Id]
356 -> Type -- Don't use; called with `undefined' below
358 dsPArrComp qs body _ =
359 dsLookupGlobalId replicatePName `thenDs` \repP ->
360 let unitArray = mkApps (Var repP) [Type unitTy,
364 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 = filterP (\x -> case x of {p -> True; _ -> False}) e
396 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
398 dePArrComp (BindStmt p e _ _ : qs) body pa cea =
399 dsLookupGlobalId filterPName `thenDs` \filterP ->
400 dsLookupGlobalId crossPName `thenDs` \crossP ->
401 dsLExpr e `thenDs` \ce ->
402 let ty'cea = parrElemType cea
403 ty'ce = parrElemType ce
404 false = Var falseDataConId
405 true = Var trueDataConId
407 newSysLocalDs ty'ce `thenDs` \v ->
408 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
409 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
410 ty'cef = ty'ce -- filterP preserves the type
411 pa' = mkTuplePat [pa, p]
413 dePArrComp qs body pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
415 -- <<[:e' | let ds, qs:]>> pa ea =
416 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
417 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
419 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
421 dePArrComp (LetStmt ds : qs) body pa cea =
422 dsLookupGlobalId mapPName `thenDs` \mapP ->
423 let xs = map unLoc (collectLocalBinders ds)
424 ty'cea = parrElemType cea
426 newSysLocalDs ty'cea `thenDs` \v ->
427 dsLocalBinds ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
428 newSysLocalDs (exprType clet) `thenDs` \let'v ->
429 let projBody = mkDsLet (NonRec let'v clet) $
430 mkCoreTup [Var v, Var let'v]
431 errTy = exprType projBody
432 errMsg = "DsListComp.dePArrComp: internal error!"
434 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
435 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase ->
436 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
437 proj = mkLams [v] ccase
439 dePArrComp qs body pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
441 -- <<[:e' | qs | qss:]>> pa ea =
442 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
443 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
445 -- {x_1, ..., x_n} = DV (qs)
447 dePArrComp (ParStmt qss : qs) body pa cea =
448 dsLookupGlobalId crossPName `thenDs` \crossP ->
449 deParStmt qss `thenDs` \(pQss,
451 let ty'cea = parrElemType cea
452 ty'ceQss = parrElemType ceQss
453 pa' = mkTuplePat [pa, pQss]
455 dePArrComp qs body pa' (mkApps (Var crossP) [Type ty'cea, Type ty'ceQss,
459 -- empty parallel statement lists have not source representation
460 panic "DsListComp.dePArrComp: Empty parallel list comprehension"
461 deParStmt ((qs, xs):qss) = -- first statement
462 let res_expr = mkExplicitTuple (map nlHsVar xs)
464 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
465 parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
467 parStmts [] pa cea = return (pa, cea)
468 parStmts ((qs, xs):qss) pa cea = -- subsequent statements (zip'ed)
469 dsLookupGlobalId zipPName `thenDs` \zipP ->
470 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
471 ty'cea = parrElemType cea
472 res_expr = mkExplicitTuple (map nlHsVar xs)
474 dsPArrComp (map unLoc qs) res_expr undefined `thenDs` \cqs ->
475 let ty'cqs = parrElemType cqs
476 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
478 parStmts qss pa' cea'
480 -- generate Core corresponding to `\p -> e'
482 deLambda :: Type -- type of the argument
483 -> LPat Id -- argument pattern
484 -> LHsExpr Id -- body
485 -> DsM (CoreExpr, Type)
487 newSysLocalDs ty `thenDs` \v ->
488 dsLExpr e `thenDs` \ce ->
489 let errTy = exprType ce
490 errMsg = "DsListComp.deLambda: internal error!"
492 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
493 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
494 returnDs (mkLams [v] res, errTy)
496 -- obtain the element type of the parallel array produced by the given Core
499 parrElemType :: CoreExpr -> Type
501 case splitTyConApp_maybe (exprType e) of
502 Just (tycon, [ty]) | tycon == parrTyCon -> ty
504 "DsListComp.parrElemType: not a parallel array type"
506 -- Smart constructor for source tuple patterns
508 mkTuplePat :: [LPat id] -> LPat id
509 mkTuplePat [lpat] = lpat
510 mkTuplePat lpats = noLoc $ TuplePat lpats Boxed
512 -- Smart constructor for source tuple expressions
514 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
515 mkExplicitTuple [lexp] = lexp
516 mkExplicitTuple lexps = noLoc $ ExplicitTuple lexps Boxed