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_FoldrBuildOn )
26 import CoreUtils ( exprType, mkIfThenElse )
29 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type,
31 import TysPrim ( alphaTyVar )
32 import TysWiredIn ( nilDataCon, consDataCon, unitDataConId, unitTy,
34 import Match ( matchSimply )
35 import PrelNames ( trueDataConName, falseDataConName, foldrName,
36 buildName, replicatePName, mapPName, filterPName,
37 zipPName, crossPName, parrTyConName )
38 import PrelInfo ( pAT_ERROR_ID )
39 import SrcLoc ( noSrcLoc )
40 import Panic ( panic )
43 List comprehensions may be desugared in one of two ways: ``ordinary''
44 (as you would expect if you read SLPJ's book) and ``with foldr/build
45 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
47 There will be at least one ``qualifier'' in the input.
50 dsListComp :: [TypecheckedStmt]
51 -> Type -- Type of list elements
54 dsListComp quals elt_ty
55 | not opt_FoldrBuildOn -- Be boring
56 || isParallelComp quals
57 = deListComp quals (mkNilExpr elt_ty)
59 | otherwise -- foldr/build lives!
60 = newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
62 n_ty = mkTyVarTy n_tyvar
63 c_ty = mkFunTys [elt_ty, n_ty] n_ty
65 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
66 dfListComp c n quals `thenDs` \ result ->
67 dsLookupGlobalId buildName `thenDs` \ build_id ->
68 returnDs (Var build_id `App` Type elt_ty
69 `App` mkLams [n_tyvar, c, n] result)
71 where isParallelComp (ParStmtOut bndrstmtss : _) = True
72 isParallelComp _ = False
75 %************************************************************************
77 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
79 %************************************************************************
81 Just as in Phil's chapter~7 in SLPJ, using the rules for
82 optimally-compiled list comprehensions. This is what Kevin followed
83 as well, and I quite happily do the same. The TQ translation scheme
84 transforms a list of qualifiers (either boolean expressions or
85 generators) into a single expression which implements the list
86 comprehension. Because we are generating 2nd-order polymorphic
87 lambda-calculus, calls to NIL and CONS must be applied to a type
88 argument, as well as their usual value arguments.
90 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
93 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
96 TQ << [ e | b , qs ] ++ L >> =
97 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
100 TQ << [ e | p <- L1, qs ] ++ L2 >> =
106 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
111 "h", "u1", "u2", and "u3" are new variables.
114 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
115 is the TE translation scheme. Note that we carry around the @L@ list
116 already desugared. @dsListComp@ does the top TE rule mentioned above.
118 To the above, we add an additional rule to deal with parallel list
119 comprehensions. The translation goes roughly as follows:
120 [ e | p1 <- e11, let v1 = e12, p2 <- e13
121 | q1 <- e21, let v2 = e22, q2 <- e23]
123 [ e | ((x1, .., xn), (y1, ..., ym)) <-
124 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
125 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
126 where (x1, .., xn) are the variables bound in p1, v1, p2
127 (y1, .., ym) are the variables bound in q1, v2, q2
129 In the translation below, the ParStmtOut branch translates each parallel branch
130 into a sub-comprehension, and desugars each independently. The resulting lists
131 are fed to a zip function, we create a binding for all the variables bound in all
132 the comprehensions, and then we hand things off the the desugarer for bindings.
133 The zip function is generated here a) because it's small, and b) because then we
134 don't have to deal with arbitrary limits on the number of zip functions in the
135 prelude, nor which library the zip function came from.
136 The introduced tuples are Boxed, but only because I couldn't get it to work
137 with the Unboxed variety.
141 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
143 deListComp (ParStmtOut bndrstmtss : quals) list
144 = mapDs do_list_comp bndrstmtss `thenDs` \ exps ->
145 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
147 -- Deal with [e | pat <- zip l1 .. ln] in example above
148 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
151 where -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
152 pat = TuplePat pats Boxed
153 pats = map (\(bs,_) -> mk_hs_tuple_pat bs) bndrstmtss
155 -- Types of (x1,..,xn), (y1,..,yn) etc
156 qual_tys = [ mk_bndrs_tys bndrs | (bndrs,_) <- bndrstmtss ]
158 do_list_comp (bndrs, stmts)
159 = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
162 mk_bndrs_tys bndrs = mk_tuple_ty (map idType bndrs)
164 -- Last: the one to return
165 deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
166 = dsExpr expr `thenDs` \ core_expr ->
167 returnDs (mkConsExpr (exprType core_expr) core_expr list)
169 -- Non-last: must be a guard
170 deListComp (ExprStmt guard ty locn : quals) list -- rule B above
171 = dsExpr guard `thenDs` \ core_guard ->
172 deListComp quals list `thenDs` \ core_rest ->
173 returnDs (mkIfThenElse core_guard core_rest list)
175 -- [e | let B, qs] = let B in [e | qs]
176 deListComp (LetStmt binds : quals) list
177 = deListComp quals list `thenDs` \ core_rest ->
178 dsLet binds core_rest
180 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
181 = dsExpr list1 `thenDs` \ core_list1 ->
182 deBindComp pat core_list1 quals core_list2
187 deBindComp pat core_list1 quals core_list2
189 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
191 -- u1_ty is a [alpha] type, and u2_ty = alpha
192 u2_ty = hsPatType pat
194 res_ty = exprType core_list2
195 h_ty = u1_ty `mkFunTy` res_ty
197 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
199 -- the "fail" value ...
201 core_fail = App (Var h) (Var u3)
202 letrec_body = App (Var h) core_list1
204 deListComp quals core_fail `thenDs` \ rest_expr ->
205 matchSimply (Var u2) (StmtCtxt ListComp) pat
206 rest_expr core_fail `thenDs` \ core_match ->
209 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
210 (DataAlt consDataCon, [u2, u3], core_match)]
212 returnDs (Let (Rec [(h, rhs)]) letrec_body)
217 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
218 -- mkZipBind [t1, t2]
219 -- = (zip, \as1:[t1] as2:[t2]
222 -- (a1:as'1) -> case as2 of
224 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
227 = mapDs newSysLocalDs list_tys `thenDs` \ ass ->
228 mapDs newSysLocalDs elt_tys `thenDs` \ as' ->
229 mapDs newSysLocalDs list_tys `thenDs` \ as's ->
230 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
232 inner_rhs = mkConsExpr ret_elt_ty (mkTupleExpr as') (mkVarApps (Var zip_fn) as's)
233 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
235 returnDs (zip_fn, mkLams ass zip_body)
237 list_tys = map mkListTy elt_tys
238 ret_elt_ty = mk_tuple_ty elt_tys
239 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
241 mk_case (as, a', as') rest
242 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
243 (DataAlt consDataCon, [a', as'], rest)]
246 mk_tuple_ty :: [Type] -> Type
247 mk_tuple_ty [ty] = ty
248 mk_tuple_ty tys = mkTupleTy Boxed (length tys) tys
250 -- Helper functions that makes an HsTuple only for non-1-sized tuples
251 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
252 mk_hs_tuple_expr [] = HsVar unitDataConId
253 mk_hs_tuple_expr [id] = HsVar id
254 mk_hs_tuple_expr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
256 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
257 mk_hs_tuple_pat [b] = VarPat b
258 mk_hs_tuple_pat bs = TuplePat (map VarPat bs) Boxed
262 %************************************************************************
264 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
266 %************************************************************************
268 @dfListComp@ are the rules used with foldr/build turned on:
271 TE[ e | ] c n = c e n
272 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
273 TE[ e | p <- l , q ] c n = let
274 f = \ x b -> case x of
282 dfListComp :: Id -> Id -- 'c' and 'n'
283 -> [TypecheckedStmt] -- the rest of the qual's
286 -- Last: the one to return
287 dfListComp c_id n_id [ResultStmt expr locn]
288 = dsExpr expr `thenDs` \ core_expr ->
289 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
291 -- Non-last: must be a guard
292 dfListComp c_id n_id (ExprStmt guard ty locn : quals)
293 = dsExpr guard `thenDs` \ core_guard ->
294 dfListComp c_id n_id quals `thenDs` \ core_rest ->
295 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
297 dfListComp c_id n_id (LetStmt binds : quals)
298 -- new in 1.3, local bindings
299 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
300 dsLet binds core_rest
302 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
303 -- evaluate the two lists
304 = dsExpr list1 `thenDs` \ core_list1 ->
306 -- find the required type
307 let x_ty = hsPatType pat
311 -- create some new local id's
312 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
314 -- build rest of the comprehesion
315 dfListComp c_id b quals `thenDs` \ core_rest ->
317 -- build the pattern match
318 matchSimply (Var x) (StmtCtxt ListComp)
319 pat core_rest (Var b) `thenDs` \ core_expr ->
321 -- now build the outermost foldr, and return
322 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
324 Var foldr_id `App` Type x_ty
326 `App` mkLams [x, b] core_expr
332 %************************************************************************
334 \subsection[DsPArrComp]{Desugaring of array comprehensions}
336 %************************************************************************
340 -- entry point for desugaring a parallel array comprehension
342 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
344 dsPArrComp :: [TypecheckedStmt]
345 -> Type -- Don't use; called with `undefined' below
348 dsLookupGlobalId replicatePName `thenDs` \repP ->
349 let unitArray = mkApps (Var repP) [Type unitTy,
353 dePArrComp qs (TuplePat [] Boxed) unitArray
357 dePArrComp :: [TypecheckedStmt]
358 -> TypecheckedPat -- the current generator pattern
359 -> CoreExpr -- the current generator expression
362 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
364 dePArrComp [ResultStmt e' _] pa cea =
365 dsLookupGlobalId mapPName `thenDs` \mapP ->
366 let ty = parrElemType cea
368 deLambda ty pa e' `thenDs` \(clam,
370 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
372 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
374 dePArrComp (ExprStmt b _ _ : qs) pa cea =
375 dsLookupGlobalId filterPName `thenDs` \filterP ->
376 let ty = parrElemType cea
378 deLambda ty pa b `thenDs` \(clam,_) ->
379 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
381 -- <<[:e' | p <- e, qs:]>> pa ea =
382 -- let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
384 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
386 dePArrComp (BindStmt p e _ : qs) pa cea =
387 dsLookupGlobalId falseDataConName `thenDs` \falseId ->
388 dsLookupGlobalId trueDataConName `thenDs` \trueId ->
389 dsLookupGlobalId filterPName `thenDs` \filterP ->
390 dsLookupGlobalId crossPName `thenDs` \crossP ->
391 dsExpr e `thenDs` \ce ->
392 let ty'cea = parrElemType cea
393 ty'ce = parrElemType ce
397 newSysLocalDs ty'ce `thenDs` \v ->
398 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
399 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
400 ty'cef = ty'ce -- filterP preserves the type
401 pa' = TuplePat [pa, p] Boxed
403 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
405 -- <<[:e' | let ds, qs:]>> pa ea =
406 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
407 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
409 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
411 dePArrComp (LetStmt ds : qs) pa cea =
412 dsLookupGlobalId mapPName `thenDs` \mapP ->
413 let xs = collectHsBinders ds
414 ty'cea = parrElemType cea
416 newSysLocalDs ty'cea `thenDs` \v ->
417 dsLet ds (mkTupleExpr xs) `thenDs` \clet ->
418 newSysLocalDs (exprType clet) `thenDs` \let'v ->
419 let projBody = mkDsLet (NonRec let'v clet) $ mkTupleExpr [v, let'v]
420 errTy = exprType projBody
421 errMsg = "DsListComp.dePArrComp: internal error!"
423 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
424 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr `thenDs` \ccase ->
425 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
426 proj = mkLams [v] ccase
428 dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
430 -- <<[:e' | qs | qss:]>> pa ea =
431 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
432 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
434 -- {x_1, ..., x_n} = DV (qs)
436 dePArrComp (ParStmtOut [] : qss2) pa cea = dePArrComp qss2 pa cea
437 dePArrComp (ParStmtOut ((xs, qs):qss) : qss2) pa cea =
438 dsLookupGlobalId zipPName `thenDs` \zipP ->
439 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
440 ty'cea = parrElemType cea
441 resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
443 dsPArrComp (qs ++ [resStmt]) undefined `thenDs` \cqs ->
444 let ty'cqs = parrElemType cqs
445 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
447 dePArrComp (ParStmtOut qss : qss2) pa' cea'
449 -- generate Core corresponding to `\p -> e'
451 deLambda :: Type -- type of the argument
452 -> TypecheckedPat -- argument pattern
453 -> TypecheckedHsExpr -- body
454 -> DsM (CoreExpr, Type)
456 newSysLocalDs ty `thenDs` \v ->
457 dsExpr e `thenDs` \ce ->
458 let errTy = exprType ce
459 errMsg = "DsListComp.deLambda: internal error!"
461 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
462 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
463 returnDs (mkLams [v] res, errTy)
465 -- obtain the element type of the parallel array produced by the given Core
468 parrElemType :: CoreExpr -> Type
470 case splitTyConApp_maybe (exprType e) of
471 Just (tycon, [ty]) | tyConName tycon == parrTyConName -> ty
473 "DsListComp.parrElemType: not a parallel array type"