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, dsLet )
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]
47 -> Type -- Type of list elements
49 dsListComp lquals elt_ty
50 = getDOptsDs `thenDs` \dflags ->
52 quals = map unLoc lquals
54 if opt_RulesOff || dopt Opt_IgnoreInterfacePragmas dflags
55 -- Either rules are switched off, or we are ignoring what there are;
56 -- Either way foldr/build won't happen, so use the more efficient
57 -- Wadler-style desugaring
58 || isParallelComp quals
59 -- Foldr-style desugaring can't handle
60 -- parallel list comprehensions
61 then deListComp quals (mkNilExpr elt_ty)
63 else -- Foldr/build should be enabled, so desugar
64 -- into foldrs and builds
65 newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
67 n_ty = mkTyVarTy n_tyvar
68 c_ty = mkFunTys [elt_ty, n_ty] n_ty
70 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
71 dfListComp c n quals `thenDs` \ result ->
72 dsLookupGlobalId buildName `thenDs` \ build_id ->
73 returnDs (Var build_id `App` Type elt_ty
74 `App` mkLams [n_tyvar, c, n] result)
76 where isParallelComp (ParStmt bndrstmtss : _) = True
77 isParallelComp _ = False
80 %************************************************************************
82 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
84 %************************************************************************
86 Just as in Phil's chapter~7 in SLPJ, using the rules for
87 optimally-compiled list comprehensions. This is what Kevin followed
88 as well, and I quite happily do the same. The TQ translation scheme
89 transforms a list of qualifiers (either boolean expressions or
90 generators) into a single expression which implements the list
91 comprehension. Because we are generating 2nd-order polymorphic
92 lambda-calculus, calls to NIL and CONS must be applied to a type
93 argument, as well as their usual value arguments.
95 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
98 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
101 TQ << [ e | b , qs ] ++ L >> =
102 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
105 TQ << [ e | p <- L1, qs ] ++ L2 >> =
111 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
116 "h", "u1", "u2", and "u3" are new variables.
119 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
120 is the TE translation scheme. Note that we carry around the @L@ list
121 already desugared. @dsListComp@ does the top TE rule mentioned above.
123 To the above, we add an additional rule to deal with parallel list
124 comprehensions. The translation goes roughly as follows:
125 [ e | p1 <- e11, let v1 = e12, p2 <- e13
126 | q1 <- e21, let v2 = e22, q2 <- e23]
128 [ e | ((x1, .., xn), (y1, ..., ym)) <-
129 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
130 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
131 where (x1, .., xn) are the variables bound in p1, v1, p2
132 (y1, .., ym) are the variables bound in q1, v2, q2
134 In the translation below, the ParStmt branch translates each parallel branch
135 into a sub-comprehension, and desugars each independently. The resulting lists
136 are fed to a zip function, we create a binding for all the variables bound in all
137 the comprehensions, and then we hand things off the the desugarer for bindings.
138 The zip function is generated here a) because it's small, and b) because then we
139 don't have to deal with arbitrary limits on the number of zip functions in the
140 prelude, nor which library the zip function came from.
141 The introduced tuples are Boxed, but only because I couldn't get it to work
142 with the Unboxed variety.
145 deListComp :: [Stmt Id] -> CoreExpr -> DsM CoreExpr
147 deListComp (ParStmt stmtss_w_bndrs : quals) list
148 = mappM do_list_comp stmtss_w_bndrs `thenDs` \ exps ->
149 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
151 -- Deal with [e | pat <- zip l1 .. ln] in example above
152 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
156 bndrs_s = map snd stmtss_w_bndrs
158 -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
159 pat = noLoc (TuplePat pats Boxed)
160 pats = map mk_hs_tuple_pat bndrs_s
162 -- Types of (x1,..,xn), (y1,..,yn) etc
163 qual_tys = map mk_bndrs_tys bndrs_s
165 do_list_comp (stmts, bndrs)
166 = dsListComp (stmts ++ [noLoc $ ResultStmt (mk_hs_tuple_expr bndrs)])
169 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
171 -- Last: the one to return
172 deListComp [ResultStmt expr] list -- Figure 7.4, SLPJ, p 135, rule C above
173 = dsLExpr expr `thenDs` \ core_expr ->
174 returnDs (mkConsExpr (exprType core_expr) core_expr list)
176 -- Non-last: must be a guard
177 deListComp (ExprStmt guard ty : quals) list -- rule B above
178 = dsLExpr guard `thenDs` \ core_guard ->
179 deListComp quals list `thenDs` \ core_rest ->
180 returnDs (mkIfThenElse core_guard core_rest list)
182 -- [e | let B, qs] = let B in [e | qs]
183 deListComp (LetStmt binds : quals) list
184 = deListComp quals list `thenDs` \ core_rest ->
185 dsLet binds core_rest
187 deListComp (BindStmt pat list1 : quals) core_list2 -- rule A' above
188 = dsLExpr list1 `thenDs` \ core_list1 ->
189 deBindComp pat core_list1 quals core_list2
194 deBindComp pat core_list1 quals core_list2
196 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
198 -- u1_ty is a [alpha] type, and u2_ty = alpha
199 u2_ty = hsPatType pat
201 res_ty = exprType core_list2
202 h_ty = u1_ty `mkFunTy` res_ty
204 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
206 -- the "fail" value ...
208 core_fail = App (Var h) (Var u3)
209 letrec_body = App (Var h) core_list1
211 deListComp quals core_fail `thenDs` \ rest_expr ->
212 matchSimply (Var u2) (StmtCtxt ListComp) pat
213 rest_expr core_fail `thenDs` \ core_match ->
216 Case (Var u1) u1 res_ty
217 [(DataAlt nilDataCon, [], core_list2),
218 (DataAlt consDataCon, [u2, u3], core_match)]
219 -- Increasing order of tag
221 returnDs (Let (Rec [(h, rhs)]) letrec_body)
226 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
227 -- mkZipBind [t1, t2]
228 -- = (zip, \as1:[t1] as2:[t2]
231 -- (a1:as'1) -> case as2 of
233 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
236 = mappM newSysLocalDs list_tys `thenDs` \ ass ->
237 mappM newSysLocalDs elt_tys `thenDs` \ as' ->
238 mappM newSysLocalDs list_tys `thenDs` \ as's ->
239 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
241 inner_rhs = mkConsExpr ret_elt_ty
242 (mkCoreTup (map Var as'))
243 (mkVarApps (Var zip_fn) as's)
244 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
246 returnDs (zip_fn, mkLams ass zip_body)
248 list_tys = map mkListTy elt_tys
249 ret_elt_ty = mkCoreTupTy elt_tys
250 list_ret_ty = mkListTy ret_elt_ty
251 zip_fn_ty = mkFunTys list_tys list_ret_ty
253 mk_case (as, a', as') rest
254 = Case (Var as) as list_ret_ty
255 [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
256 (DataAlt consDataCon, [a', as'], rest)]
257 -- Increasing order of tag
258 -- Helper functions that makes an HsTuple only for non-1-sized tuples
259 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
260 mk_hs_tuple_expr [] = nlHsVar unitDataConId
261 mk_hs_tuple_expr [id] = nlHsVar id
262 mk_hs_tuple_expr ids = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
264 mk_hs_tuple_pat :: [Id] -> LPat Id
265 mk_hs_tuple_pat [b] = nlVarPat b
266 mk_hs_tuple_pat bs = noLoc $ TuplePat (map nlVarPat bs) Boxed
270 %************************************************************************
272 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
274 %************************************************************************
276 @dfListComp@ are the rules used with foldr/build turned on:
279 TE[ e | ] c n = c e n
280 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
281 TE[ e | p <- l , q ] c n = let
282 f = \ x b -> case x of
290 dfListComp :: Id -> Id -- 'c' and 'n'
291 -> [Stmt Id] -- the rest of the qual's
294 -- Last: the one to return
295 dfListComp c_id n_id [ResultStmt expr]
296 = dsLExpr expr `thenDs` \ core_expr ->
297 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
299 -- Non-last: must be a guard
300 dfListComp c_id n_id (ExprStmt guard ty : quals)
301 = dsLExpr guard `thenDs` \ core_guard ->
302 dfListComp c_id n_id quals `thenDs` \ core_rest ->
303 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
305 dfListComp c_id n_id (LetStmt binds : quals)
306 -- new in 1.3, local bindings
307 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
308 dsLet binds core_rest
310 dfListComp c_id n_id (BindStmt pat list1 : quals)
311 -- evaluate the two lists
312 = dsLExpr list1 `thenDs` \ core_list1 ->
314 -- find the required type
315 let x_ty = hsPatType pat
319 -- create some new local id's
320 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
322 -- build rest of the comprehesion
323 dfListComp c_id b quals `thenDs` \ core_rest ->
325 -- build the pattern match
326 matchSimply (Var x) (StmtCtxt ListComp)
327 pat core_rest (Var b) `thenDs` \ core_expr ->
329 -- now build the outermost foldr, and return
330 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
332 Var foldr_id `App` Type x_ty
334 `App` mkLams [x, b] core_expr
340 %************************************************************************
342 \subsection[DsPArrComp]{Desugaring of array comprehensions}
344 %************************************************************************
348 -- entry point for desugaring a parallel array comprehension
350 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
352 dsPArrComp :: [Stmt Id]
353 -> Type -- Don't use; called with `undefined' below
356 dsLookupGlobalId replicatePName `thenDs` \repP ->
357 let unitArray = mkApps (Var repP) [Type unitTy,
361 dePArrComp qs (mkTuplePat []) unitArray
365 dePArrComp :: [Stmt Id]
366 -> LPat Id -- the current generator pattern
367 -> CoreExpr -- the current generator expression
370 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
372 dePArrComp [ResultStmt e'] pa cea =
373 dsLookupGlobalId mapPName `thenDs` \mapP ->
374 let ty = parrElemType cea
376 deLambda ty pa e' `thenDs` \(clam,
378 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
380 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
382 dePArrComp (ExprStmt b _ : qs) pa cea =
383 dsLookupGlobalId filterPName `thenDs` \filterP ->
384 let ty = parrElemType cea
386 deLambda ty pa b `thenDs` \(clam,_) ->
387 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
389 -- <<[:e' | p <- e, qs:]>> pa ea =
390 -- let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
392 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
394 dePArrComp (BindStmt p e : qs) pa cea =
395 dsLookupGlobalId filterPName `thenDs` \filterP ->
396 dsLookupGlobalId crossPName `thenDs` \crossP ->
397 dsLExpr e `thenDs` \ce ->
398 let ty'cea = parrElemType cea
399 ty'ce = parrElemType ce
400 false = Var falseDataConId
401 true = Var trueDataConId
403 newSysLocalDs ty'ce `thenDs` \v ->
404 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
405 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
406 ty'cef = ty'ce -- filterP preserves the type
407 pa' = mkTuplePat [pa, p]
409 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
411 -- <<[:e' | let ds, qs:]>> pa ea =
412 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
413 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
415 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
417 dePArrComp (LetStmt ds : qs) pa cea =
418 dsLookupGlobalId mapPName `thenDs` \mapP ->
419 let xs = map unLoc (collectGroupBinders ds)
420 ty'cea = parrElemType cea
422 newSysLocalDs ty'cea `thenDs` \v ->
423 dsLet ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
424 newSysLocalDs (exprType clet) `thenDs` \let'v ->
425 let projBody = mkDsLet (NonRec let'v clet) $
426 mkCoreTup [Var v, Var let'v]
427 errTy = exprType projBody
428 errMsg = "DsListComp.dePArrComp: internal error!"
430 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
431 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase ->
432 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
433 proj = mkLams [v] ccase
435 dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
437 -- <<[:e' | qs | qss:]>> pa ea =
438 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
439 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
441 -- {x_1, ..., x_n} = DV (qs)
443 dePArrComp (ParStmt qss : qs) pa cea =
444 dsLookupGlobalId crossPName `thenDs` \crossP ->
445 deParStmt qss `thenDs` \(pQss,
447 let ty'cea = parrElemType cea
448 ty'ceQss = parrElemType ceQss
449 pa' = mkTuplePat [pa, pQss]
451 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'ceQss,
455 -- empty parallel statement lists have not source representation
456 panic "DsListComp.dePArrComp: Empty parallel list comprehension"
457 deParStmt ((qs, xs):qss) = -- first statement
458 let resStmt = ResultStmt $ mkExplicitTuple (map nlHsVar xs)
460 dsPArrComp (map unLoc qs ++ [resStmt]) undefined `thenDs` \cqs ->
461 parStmts qss (mkTuplePat (map nlVarPat xs)) cqs
463 parStmts [] pa cea = return (pa, cea)
464 parStmts ((qs, xs):qss) pa cea = -- subsequent statements (zip'ed)
465 dsLookupGlobalId zipPName `thenDs` \zipP ->
466 let pa' = mkTuplePat [pa, mkTuplePat (map nlVarPat xs)]
467 ty'cea = parrElemType cea
468 resStmt = ResultStmt $ mkExplicitTuple (map nlHsVar xs)
470 dsPArrComp (map unLoc qs ++ [resStmt]) undefined `thenDs` \cqs ->
471 let ty'cqs = parrElemType cqs
472 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
474 parStmts qss pa' cea'
476 -- generate Core corresponding to `\p -> e'
478 deLambda :: Type -- type of the argument
479 -> LPat Id -- argument pattern
480 -> LHsExpr Id -- body
481 -> DsM (CoreExpr, Type)
483 newSysLocalDs ty `thenDs` \v ->
484 dsLExpr e `thenDs` \ce ->
485 let errTy = exprType ce
486 errMsg = "DsListComp.deLambda: internal error!"
488 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
489 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
490 returnDs (mkLams [v] res, errTy)
492 -- obtain the element type of the parallel array produced by the given Core
495 parrElemType :: CoreExpr -> Type
497 case splitTyConApp_maybe (exprType e) of
498 Just (tycon, [ty]) | tycon == parrTyCon -> ty
500 "DsListComp.parrElemType: not a parallel array type"
502 -- Smart constructor for source tuple patterns
504 mkTuplePat :: [LPat id] -> LPat id
505 mkTuplePat [lpat] = lpat
506 mkTuplePat lpats = noLoc $ TuplePat lpats Boxed
508 -- Smart constructor for source tuple expressions
510 mkExplicitTuple :: [LHsExpr id] -> LHsExpr id
511 mkExplicitTuple [lexp] = lexp
512 mkExplicitTuple lexps = noLoc $ ExplicitTuple lexps Boxed