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 ( dsExpr, dsLet )
13 import BasicTypes ( Boxity(..) )
14 import TyCon ( tyConName )
15 import HsSyn ( Pat(..), HsExpr(..), Stmt(..),
16 HsMatchContext(..), HsStmtContext(..),
18 import TcHsSyn ( TypecheckedStmt, TypecheckedPat, TypecheckedHsExpr,
22 import DsMonad -- the monadery used in the desugarer
25 import CmdLineOpts ( opt_IgnoreIfacePragmas, opt_RulesOff )
26 import CoreUtils ( exprType, mkIfThenElse )
29 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type,
31 import TysPrim ( alphaTyVar )
32 import TysWiredIn ( nilDataCon, consDataCon, trueDataConId, falseDataConId,
33 unitDataConId, unitTy, mkListTy )
34 import Match ( matchSimply )
35 import PrelNames ( foldrName, buildName, replicatePName, mapPName,
36 filterPName, zipPName, crossPName, parrTyConName )
37 import PrelInfo ( pAT_ERROR_ID )
38 import SrcLoc ( noSrcLoc )
39 import Panic ( panic )
42 List comprehensions may be desugared in one of two ways: ``ordinary''
43 (as you would expect if you read SLPJ's book) and ``with foldr/build
44 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
46 There will be at least one ``qualifier'' in the input.
49 dsListComp :: [TypecheckedStmt]
50 -> Type -- Type of list elements
53 dsListComp quals elt_ty
54 | opt_RulesOff || opt_IgnoreIfacePragmas -- Either rules are switched off, or
55 -- we are ignoring what there are;
56 -- Either way foldr/build won't happen, so
57 -- use the more efficient Wadler-style desugaring
58 || isParallelComp quals -- Foldr-style desugaring can't handle
59 -- parallel list comprehensions
60 = deListComp quals (mkNilExpr elt_ty)
62 | otherwise -- 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.
145 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
147 deListComp (ParStmt stmtss_w_bndrs : quals) list
148 = mapDs 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 = 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 ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
169 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
171 -- Last: the one to return
172 deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
173 = dsExpr 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 locn : quals) list -- rule B above
178 = dsExpr 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 locn : quals) core_list2 -- rule A' above
188 = dsExpr 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 [(DataAlt nilDataCon, [], core_list2),
217 (DataAlt consDataCon, [u2, u3], core_match)]
219 returnDs (Let (Rec [(h, rhs)]) letrec_body)
224 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
225 -- mkZipBind [t1, t2]
226 -- = (zip, \as1:[t1] as2:[t2]
229 -- (a1:as'1) -> case as2 of
231 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
234 = mapDs newSysLocalDs list_tys `thenDs` \ ass ->
235 mapDs newSysLocalDs elt_tys `thenDs` \ as' ->
236 mapDs newSysLocalDs list_tys `thenDs` \ as's ->
237 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
239 inner_rhs = mkConsExpr ret_elt_ty
240 (mkCoreTup (map Var as'))
241 (mkVarApps (Var zip_fn) as's)
242 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
244 returnDs (zip_fn, mkLams ass zip_body)
246 list_tys = map mkListTy elt_tys
247 ret_elt_ty = mkCoreTupTy elt_tys
248 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
250 mk_case (as, a', as') rest
251 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
252 (DataAlt consDataCon, [a', as'], rest)]
254 -- Helper functions that makes an HsTuple only for non-1-sized tuples
255 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
256 mk_hs_tuple_expr [] = HsVar unitDataConId
257 mk_hs_tuple_expr [id] = HsVar id
258 mk_hs_tuple_expr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
260 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
261 mk_hs_tuple_pat [b] = VarPat b
262 mk_hs_tuple_pat bs = TuplePat (map VarPat bs) Boxed
266 %************************************************************************
268 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
270 %************************************************************************
272 @dfListComp@ are the rules used with foldr/build turned on:
275 TE[ e | ] c n = c e n
276 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
277 TE[ e | p <- l , q ] c n = let
278 f = \ x b -> case x of
286 dfListComp :: Id -> Id -- 'c' and 'n'
287 -> [TypecheckedStmt] -- the rest of the qual's
290 -- Last: the one to return
291 dfListComp c_id n_id [ResultStmt expr locn]
292 = dsExpr expr `thenDs` \ core_expr ->
293 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
295 -- Non-last: must be a guard
296 dfListComp c_id n_id (ExprStmt guard ty locn : quals)
297 = dsExpr guard `thenDs` \ core_guard ->
298 dfListComp c_id n_id quals `thenDs` \ core_rest ->
299 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
301 dfListComp c_id n_id (LetStmt binds : quals)
302 -- new in 1.3, local bindings
303 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
304 dsLet binds core_rest
306 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
307 -- evaluate the two lists
308 = dsExpr list1 `thenDs` \ core_list1 ->
310 -- find the required type
311 let x_ty = hsPatType pat
315 -- create some new local id's
316 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
318 -- build rest of the comprehesion
319 dfListComp c_id b quals `thenDs` \ core_rest ->
321 -- build the pattern match
322 matchSimply (Var x) (StmtCtxt ListComp)
323 pat core_rest (Var b) `thenDs` \ core_expr ->
325 -- now build the outermost foldr, and return
326 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
328 Var foldr_id `App` Type x_ty
330 `App` mkLams [x, b] core_expr
336 %************************************************************************
338 \subsection[DsPArrComp]{Desugaring of array comprehensions}
340 %************************************************************************
344 -- entry point for desugaring a parallel array comprehension
346 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
348 dsPArrComp :: [TypecheckedStmt]
349 -> Type -- Don't use; called with `undefined' below
352 dsLookupGlobalId replicatePName `thenDs` \repP ->
353 let unitArray = mkApps (Var repP) [Type unitTy,
357 dePArrComp qs (TuplePat [] Boxed) unitArray
361 dePArrComp :: [TypecheckedStmt]
362 -> TypecheckedPat -- the current generator pattern
363 -> CoreExpr -- the current generator expression
366 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
368 dePArrComp [ResultStmt e' _] pa cea =
369 dsLookupGlobalId mapPName `thenDs` \mapP ->
370 let ty = parrElemType cea
372 deLambda ty pa e' `thenDs` \(clam,
374 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
376 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
378 dePArrComp (ExprStmt b _ _ : qs) pa cea =
379 dsLookupGlobalId filterPName `thenDs` \filterP ->
380 let ty = parrElemType cea
382 deLambda ty pa b `thenDs` \(clam,_) ->
383 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
385 -- <<[:e' | p <- e, qs:]>> pa ea =
386 -- let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
388 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
390 dePArrComp (BindStmt p e _ : qs) pa cea =
391 dsLookupGlobalId filterPName `thenDs` \filterP ->
392 dsLookupGlobalId crossPName `thenDs` \crossP ->
393 dsExpr e `thenDs` \ce ->
394 let ty'cea = parrElemType cea
395 ty'ce = parrElemType ce
396 false = Var falseDataConId
397 true = Var trueDataConId
399 newSysLocalDs ty'ce `thenDs` \v ->
400 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
401 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
402 ty'cef = ty'ce -- filterP preserves the type
403 pa' = TuplePat [pa, p] Boxed
405 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
407 -- <<[:e' | let ds, qs:]>> pa ea =
408 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
409 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
411 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
413 dePArrComp (LetStmt ds : qs) pa cea =
414 dsLookupGlobalId mapPName `thenDs` \mapP ->
415 let xs = collectHsBinders ds
416 ty'cea = parrElemType cea
418 newSysLocalDs ty'cea `thenDs` \v ->
419 dsLet ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
420 newSysLocalDs (exprType clet) `thenDs` \let'v ->
421 let projBody = mkDsLet (NonRec let'v clet) $
422 mkCoreTup [Var v, Var let'v]
423 errTy = exprType projBody
424 errMsg = "DsListComp.dePArrComp: internal error!"
426 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
427 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr `thenDs` \ccase ->
428 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
429 proj = mkLams [v] ccase
431 dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
433 -- <<[:e' | qs | qss:]>> pa ea =
434 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
435 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
437 -- {x_1, ..., x_n} = DV (qs)
439 dePArrComp (ParStmt [] : qss2) pa cea = dePArrComp qss2 pa cea
440 dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
441 dsLookupGlobalId zipPName `thenDs` \zipP ->
442 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
443 ty'cea = parrElemType cea
444 resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
446 dsPArrComp (qs ++ [resStmt]) undefined `thenDs` \cqs ->
447 let ty'cqs = parrElemType cqs
448 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
450 dePArrComp (ParStmt qss : qss2) pa' cea'
452 -- generate Core corresponding to `\p -> e'
454 deLambda :: Type -- type of the argument
455 -> TypecheckedPat -- argument pattern
456 -> TypecheckedHsExpr -- body
457 -> DsM (CoreExpr, Type)
459 newSysLocalDs ty `thenDs` \v ->
460 dsExpr e `thenDs` \ce ->
461 let errTy = exprType ce
462 errMsg = "DsListComp.deLambda: internal error!"
464 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
465 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
466 returnDs (mkLams [v] res, errTy)
468 -- obtain the element type of the parallel array produced by the given Core
471 parrElemType :: CoreExpr -> Type
473 case splitTyConApp_maybe (exprType e) of
474 Just (tycon, [ty]) | tyConName tycon == parrTyConName -> ty
476 "DsListComp.parrElemType: not a parallel array type"