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 CmdLineOpts ( DynFlag(..), dopt, opt_RulesOff )
22 import CoreUtils ( exprType, mkIfThenElse )
25 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type,
27 import TysPrim ( alphaTyVar )
28 import TysWiredIn ( nilDataCon, consDataCon, trueDataConId, falseDataConId,
29 unitDataConId, unitTy, mkListTy, parrTyCon )
30 import Match ( matchSimply )
31 import PrelNames ( foldrName, buildName, replicatePName, mapPName,
32 filterPName, zipPName, crossPName )
33 import PrelInfo ( pAT_ERROR_ID )
34 import SrcLoc ( noLoc, unLoc )
35 import Panic ( panic )
38 List comprehensions may be desugared in one of two ways: ``ordinary''
39 (as you would expect if you read SLPJ's book) and ``with foldr/build
40 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
42 There will be at least one ``qualifier'' in the input.
45 dsListComp :: [LStmt Id]
46 -> Type -- Type of list elements
48 dsListComp lquals 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 (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 `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] -> CoreExpr -> DsM CoreExpr
146 deListComp (ParStmt stmtss_w_bndrs : quals) 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 = noLoc (TuplePat pats Boxed)
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 ++ [noLoc $ ResultStmt (mk_hs_tuple_expr bndrs)])
168 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
170 -- Last: the one to return
171 deListComp [ResultStmt expr] list -- Figure 7.4, SLPJ, p 135, rule C above
172 = dsLExpr expr `thenDs` \ core_expr ->
173 returnDs (mkConsExpr (exprType core_expr) core_expr list)
175 -- Non-last: must be a guard
176 deListComp (ExprStmt guard ty : quals) list -- rule B above
177 = dsLExpr guard `thenDs` \ core_guard ->
178 deListComp quals 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) list
183 = deListComp quals list `thenDs` \ core_rest ->
184 dsLet binds core_rest
186 deListComp (BindStmt pat list1 : quals) core_list2 -- rule A' above
187 = dsLExpr list1 `thenDs` \ core_list1 ->
188 deBindComp pat core_list1 quals core_list2
193 deBindComp pat core_list1 quals 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 = hsPatType 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 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 [(DataAlt nilDataCon, [], core_list2),
216 (DataAlt consDataCon, [u2, u3], core_match)]
218 returnDs (Let (Rec [(h, rhs)]) letrec_body)
223 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
224 -- mkZipBind [t1, t2]
225 -- = (zip, \as1:[t1] as2:[t2]
228 -- (a1:as'1) -> case as2 of
230 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
233 = mappM newSysLocalDs list_tys `thenDs` \ ass ->
234 mappM newSysLocalDs elt_tys `thenDs` \ as' ->
235 mappM newSysLocalDs list_tys `thenDs` \ as's ->
236 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
238 inner_rhs = mkConsExpr ret_elt_ty
239 (mkCoreTup (map Var as'))
240 (mkVarApps (Var zip_fn) as's)
241 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
243 returnDs (zip_fn, mkLams ass zip_body)
245 list_tys = map mkListTy elt_tys
246 ret_elt_ty = mkCoreTupTy elt_tys
247 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
249 mk_case (as, a', as') rest
250 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
251 (DataAlt consDataCon, [a', as'], rest)]
253 -- Helper functions that makes an HsTuple only for non-1-sized tuples
254 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
255 mk_hs_tuple_expr [] = nlHsVar unitDataConId
256 mk_hs_tuple_expr [id] = nlHsVar id
257 mk_hs_tuple_expr ids = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
259 mk_hs_tuple_pat :: [Id] -> LPat Id
260 mk_hs_tuple_pat [b] = nlVarPat b
261 mk_hs_tuple_pat bs = noLoc $ TuplePat (map nlVarPat bs) Boxed
265 %************************************************************************
267 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
269 %************************************************************************
271 @dfListComp@ are the rules used with foldr/build turned on:
274 TE[ e | ] c n = c e n
275 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
276 TE[ e | p <- l , q ] c n = let
277 f = \ x b -> case x of
285 dfListComp :: Id -> Id -- 'c' and 'n'
286 -> [Stmt Id] -- the rest of the qual's
289 -- Last: the one to return
290 dfListComp c_id n_id [ResultStmt expr]
291 = dsLExpr expr `thenDs` \ core_expr ->
292 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
294 -- Non-last: must be a guard
295 dfListComp c_id n_id (ExprStmt guard ty : quals)
296 = dsLExpr guard `thenDs` \ core_guard ->
297 dfListComp c_id n_id quals `thenDs` \ core_rest ->
298 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
300 dfListComp c_id n_id (LetStmt binds : quals)
301 -- new in 1.3, local bindings
302 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
303 dsLet binds core_rest
305 dfListComp c_id n_id (BindStmt pat list1 : quals)
306 -- evaluate the two lists
307 = dsLExpr list1 `thenDs` \ core_list1 ->
309 -- find the required type
310 let x_ty = hsPatType pat
314 -- create some new local id's
315 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
317 -- build rest of the comprehesion
318 dfListComp c_id b quals `thenDs` \ core_rest ->
320 -- build the pattern match
321 matchSimply (Var x) (StmtCtxt ListComp)
322 pat core_rest (Var b) `thenDs` \ core_expr ->
324 -- now build the outermost foldr, and return
325 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
327 Var foldr_id `App` Type x_ty
329 `App` mkLams [x, b] core_expr
335 %************************************************************************
337 \subsection[DsPArrComp]{Desugaring of array comprehensions}
339 %************************************************************************
343 -- entry point for desugaring a parallel array comprehension
345 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
347 dsPArrComp :: [Stmt Id]
348 -> Type -- Don't use; called with `undefined' below
351 dsLookupGlobalId replicatePName `thenDs` \repP ->
352 let unitArray = mkApps (Var repP) [Type unitTy,
356 dePArrComp qs (noLoc (TuplePat [] Boxed)) unitArray
360 dePArrComp :: [Stmt Id]
361 -> LPat Id -- the current generator pattern
362 -> CoreExpr -- the current generator expression
365 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
367 dePArrComp [ResultStmt e'] pa cea =
368 dsLookupGlobalId mapPName `thenDs` \mapP ->
369 let ty = parrElemType cea
371 deLambda ty pa e' `thenDs` \(clam,
373 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
375 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
377 dePArrComp (ExprStmt b _ : qs) pa cea =
378 dsLookupGlobalId filterPName `thenDs` \filterP ->
379 let ty = parrElemType cea
381 deLambda ty pa b `thenDs` \(clam,_) ->
382 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
384 -- <<[:e' | p <- e, qs:]>> pa ea =
385 -- let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
387 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
389 dePArrComp (BindStmt p e : qs) pa cea =
390 dsLookupGlobalId filterPName `thenDs` \filterP ->
391 dsLookupGlobalId crossPName `thenDs` \crossP ->
392 dsLExpr e `thenDs` \ce ->
393 let ty'cea = parrElemType cea
394 ty'ce = parrElemType ce
395 false = Var falseDataConId
396 true = Var trueDataConId
398 newSysLocalDs ty'ce `thenDs` \v ->
399 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
400 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
401 ty'cef = ty'ce -- filterP preserves the type
402 pa' = noLoc (TuplePat [pa, p] Boxed)
404 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
406 -- <<[:e' | let ds, qs:]>> pa ea =
407 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
408 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
410 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
412 dePArrComp (LetStmt ds : qs) pa cea =
413 dsLookupGlobalId mapPName `thenDs` \mapP ->
414 let xs = map unLoc (collectGroupBinders ds)
415 ty'cea = parrElemType cea
417 newSysLocalDs ty'cea `thenDs` \v ->
418 dsLet ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
419 newSysLocalDs (exprType clet) `thenDs` \let'v ->
420 let projBody = mkDsLet (NonRec let'v clet) $
421 mkCoreTup [Var v, Var let'v]
422 errTy = exprType projBody
423 errMsg = "DsListComp.dePArrComp: internal error!"
425 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
426 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr `thenDs` \ccase ->
427 let pa' = noLoc $ TuplePat [pa, noLoc (TuplePat (map nlVarPat xs) Boxed)] Boxed
428 proj = mkLams [v] ccase
430 dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
432 -- <<[:e' | qs | qss:]>> pa ea =
433 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
434 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
436 -- {x_1, ..., x_n} = DV (qs)
438 dePArrComp (ParStmt [] : qss2) pa cea = dePArrComp qss2 pa cea
439 dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
440 dsLookupGlobalId zipPName `thenDs` \zipP ->
441 let pa' = noLoc $ TuplePat [pa, noLoc (TuplePat (map nlVarPat xs) Boxed)] Boxed
442 ty'cea = parrElemType cea
443 resStmt = ResultStmt (noLoc $ ExplicitTuple (map nlHsVar xs) Boxed)
445 dsPArrComp (map unLoc qs ++ [resStmt]) undefined `thenDs` \cqs ->
446 let ty'cqs = parrElemType cqs
447 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
449 dePArrComp (ParStmt qss : qss2) pa' cea'
451 -- generate Core corresponding to `\p -> e'
453 deLambda :: Type -- type of the argument
454 -> LPat Id -- argument pattern
455 -> LHsExpr Id -- body
456 -> DsM (CoreExpr, Type)
458 newSysLocalDs ty `thenDs` \v ->
459 dsLExpr e `thenDs` \ce ->
460 let errTy = exprType ce
461 errMsg = "DsListComp.deLambda: internal error!"
463 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
464 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
465 returnDs (mkLams [v] res, errTy)
467 -- obtain the element type of the parallel array produced by the given Core
470 parrElemType :: CoreExpr -> Type
472 case splitTyConApp_maybe (exprType e) of
473 Just (tycon, [ty]) | tycon == parrTyCon -> ty
475 "DsListComp.parrElemType: not a parallel array type"