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 HsSyn ( Pat(..), HsExpr(..), Stmt(..),
15 HsMatchContext(..), HsStmtContext(..),
17 import TcHsSyn ( TypecheckedStmt, TypecheckedPat, TypecheckedHsExpr,
21 import DsMonad -- the monadery used in the desugarer
24 import CmdLineOpts ( DynFlag(..), dopt, opt_RulesOff )
25 import CoreUtils ( exprType, mkIfThenElse )
28 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type,
30 import TysPrim ( alphaTyVar )
31 import TysWiredIn ( nilDataCon, consDataCon, trueDataConId, falseDataConId,
32 unitDataConId, unitTy, mkListTy, parrTyCon )
33 import Match ( matchSimply )
34 import PrelNames ( foldrName, buildName, replicatePName, mapPName,
35 filterPName, zipPName, crossPName )
36 import PrelInfo ( pAT_ERROR_ID )
37 import SrcLoc ( noSrcLoc )
38 import Panic ( panic )
41 List comprehensions may be desugared in one of two ways: ``ordinary''
42 (as you would expect if you read SLPJ's book) and ``with foldr/build
43 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
45 There will be at least one ``qualifier'' in the input.
48 dsListComp :: [TypecheckedStmt]
49 -> Type -- Type of list elements
52 dsListComp quals elt_ty
53 = getDOptsDs `thenDs` \dflags ->
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.
146 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
148 deListComp (ParStmt stmtss_w_bndrs : quals) 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 = 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 ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
170 mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
172 -- Last: the one to return
173 deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
174 = dsExpr expr `thenDs` \ core_expr ->
175 returnDs (mkConsExpr (exprType core_expr) core_expr list)
177 -- Non-last: must be a guard
178 deListComp (ExprStmt guard ty locn : quals) list -- rule B above
179 = dsExpr guard `thenDs` \ core_guard ->
180 deListComp quals 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) list
185 = deListComp quals list `thenDs` \ core_rest ->
186 dsLet binds core_rest
188 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
189 = dsExpr list1 `thenDs` \ core_list1 ->
190 deBindComp pat core_list1 quals core_list2
195 deBindComp pat core_list1 quals 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 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 [(DataAlt nilDataCon, [], core_list2),
218 (DataAlt consDataCon, [u2, u3], core_match)]
220 returnDs (Let (Rec [(h, rhs)]) letrec_body)
225 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
226 -- mkZipBind [t1, t2]
227 -- = (zip, \as1:[t1] as2:[t2]
230 -- (a1:as'1) -> case as2 of
232 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
235 = mappM newSysLocalDs list_tys `thenDs` \ ass ->
236 mappM newSysLocalDs elt_tys `thenDs` \ as' ->
237 mappM newSysLocalDs list_tys `thenDs` \ as's ->
238 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
240 inner_rhs = mkConsExpr ret_elt_ty
241 (mkCoreTup (map Var as'))
242 (mkVarApps (Var zip_fn) as's)
243 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
245 returnDs (zip_fn, mkLams ass zip_body)
247 list_tys = map mkListTy elt_tys
248 ret_elt_ty = mkCoreTupTy elt_tys
249 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
251 mk_case (as, a', as') rest
252 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
253 (DataAlt consDataCon, [a', as'], rest)]
255 -- Helper functions that makes an HsTuple only for non-1-sized tuples
256 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
257 mk_hs_tuple_expr [] = HsVar unitDataConId
258 mk_hs_tuple_expr [id] = HsVar id
259 mk_hs_tuple_expr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
261 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
262 mk_hs_tuple_pat [b] = VarPat b
263 mk_hs_tuple_pat bs = TuplePat (map VarPat bs) Boxed
267 %************************************************************************
269 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
271 %************************************************************************
273 @dfListComp@ are the rules used with foldr/build turned on:
276 TE[ e | ] c n = c e n
277 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
278 TE[ e | p <- l , q ] c n = let
279 f = \ x b -> case x of
287 dfListComp :: Id -> Id -- 'c' and 'n'
288 -> [TypecheckedStmt] -- the rest of the qual's
291 -- Last: the one to return
292 dfListComp c_id n_id [ResultStmt expr locn]
293 = dsExpr expr `thenDs` \ core_expr ->
294 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
296 -- Non-last: must be a guard
297 dfListComp c_id n_id (ExprStmt guard ty locn : quals)
298 = dsExpr guard `thenDs` \ core_guard ->
299 dfListComp c_id n_id quals `thenDs` \ core_rest ->
300 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
302 dfListComp c_id n_id (LetStmt binds : quals)
303 -- new in 1.3, local bindings
304 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
305 dsLet binds core_rest
307 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
308 -- evaluate the two lists
309 = dsExpr list1 `thenDs` \ core_list1 ->
311 -- find the required type
312 let x_ty = hsPatType pat
316 -- create some new local id's
317 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
319 -- build rest of the comprehesion
320 dfListComp c_id b quals `thenDs` \ core_rest ->
322 -- build the pattern match
323 matchSimply (Var x) (StmtCtxt ListComp)
324 pat core_rest (Var b) `thenDs` \ core_expr ->
326 -- now build the outermost foldr, and return
327 dsLookupGlobalId foldrName `thenDs` \ foldr_id ->
329 Var foldr_id `App` Type x_ty
331 `App` mkLams [x, b] core_expr
337 %************************************************************************
339 \subsection[DsPArrComp]{Desugaring of array comprehensions}
341 %************************************************************************
345 -- entry point for desugaring a parallel array comprehension
347 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
349 dsPArrComp :: [TypecheckedStmt]
350 -> Type -- Don't use; called with `undefined' below
353 dsLookupGlobalId replicatePName `thenDs` \repP ->
354 let unitArray = mkApps (Var repP) [Type unitTy,
358 dePArrComp qs (TuplePat [] Boxed) unitArray
362 dePArrComp :: [TypecheckedStmt]
363 -> TypecheckedPat -- the current generator pattern
364 -> CoreExpr -- the current generator expression
367 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
369 dePArrComp [ResultStmt e' _] pa cea =
370 dsLookupGlobalId mapPName `thenDs` \mapP ->
371 let ty = parrElemType cea
373 deLambda ty pa e' `thenDs` \(clam,
375 returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
377 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
379 dePArrComp (ExprStmt b _ _ : qs) pa cea =
380 dsLookupGlobalId filterPName `thenDs` \filterP ->
381 let ty = parrElemType cea
383 deLambda ty pa b `thenDs` \(clam,_) ->
384 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
386 -- <<[:e' | p <- e, qs:]>> pa ea =
387 -- let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
389 -- <<[:e' | qs:]>> (pa, p) (crossP ea ef)
391 dePArrComp (BindStmt p e _ : qs) pa cea =
392 dsLookupGlobalId filterPName `thenDs` \filterP ->
393 dsLookupGlobalId crossPName `thenDs` \crossP ->
394 dsExpr e `thenDs` \ce ->
395 let ty'cea = parrElemType cea
396 ty'ce = parrElemType ce
397 false = Var falseDataConId
398 true = Var trueDataConId
400 newSysLocalDs ty'ce `thenDs` \v ->
401 matchSimply (Var v) (StmtCtxt PArrComp) p true false `thenDs` \pred ->
402 let cef = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
403 ty'cef = ty'ce -- filterP preserves the type
404 pa' = TuplePat [pa, p] Boxed
406 dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
408 -- <<[:e' | let ds, qs:]>> pa ea =
409 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
410 -- (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
412 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
414 dePArrComp (LetStmt ds : qs) pa cea =
415 dsLookupGlobalId mapPName `thenDs` \mapP ->
416 let xs = collectHsBinders ds
417 ty'cea = parrElemType cea
419 newSysLocalDs ty'cea `thenDs` \v ->
420 dsLet ds (mkCoreTup (map Var xs)) `thenDs` \clet ->
421 newSysLocalDs (exprType clet) `thenDs` \let'v ->
422 let projBody = mkDsLet (NonRec let'v clet) $
423 mkCoreTup [Var v, Var let'v]
424 errTy = exprType projBody
425 errMsg = "DsListComp.dePArrComp: internal error!"
427 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
428 matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr `thenDs` \ccase ->
429 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
430 proj = mkLams [v] ccase
432 dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
434 -- <<[:e' | qs | qss:]>> pa ea =
435 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
436 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
438 -- {x_1, ..., x_n} = DV (qs)
440 dePArrComp (ParStmt [] : qss2) pa cea = dePArrComp qss2 pa cea
441 dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
442 dsLookupGlobalId zipPName `thenDs` \zipP ->
443 let pa' = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
444 ty'cea = parrElemType cea
445 resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
447 dsPArrComp (qs ++ [resStmt]) undefined `thenDs` \cqs ->
448 let ty'cqs = parrElemType cqs
449 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
451 dePArrComp (ParStmt qss : qss2) pa' cea'
453 -- generate Core corresponding to `\p -> e'
455 deLambda :: Type -- type of the argument
456 -> TypecheckedPat -- argument pattern
457 -> TypecheckedHsExpr -- body
458 -> DsM (CoreExpr, Type)
460 newSysLocalDs ty `thenDs` \v ->
461 dsExpr e `thenDs` \ce ->
462 let errTy = exprType ce
463 errMsg = "DsListComp.deLambda: internal error!"
465 mkErrorAppDs pAT_ERROR_ID errTy errMsg `thenDs` \cerr ->
466 matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr `thenDs` \res ->
467 returnDs (mkLams [v] res, errTy)
469 -- obtain the element type of the parallel array produced by the given Core
472 parrElemType :: CoreExpr -> Type
474 case splitTyConApp_maybe (exprType e) of
475 Just (tycon, [ty]) | tycon == parrTyCon -> ty
477 "DsListComp.parrElemType: not a parallel array type"