2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[DsListComp]{Desugaring list comprehensions}
7 module DsListComp ( dsListComp ) where
9 #include "HsVersions.h"
11 import {-# SOURCE #-} DsExpr ( dsExpr, dsLet )
13 import BasicTypes ( Boxity(..) )
14 import HsSyn ( OutPat(..), HsExpr(..), Stmt(..), HsMatchContext(..) )
15 import TcHsSyn ( TypecheckedStmt )
16 import DsHsSyn ( outPatType )
19 import DsMonad -- the monadery used in the desugarer
22 import CmdLineOpts ( opt_FoldrBuildOn )
23 import CoreUtils ( exprType, mkIfThenElse )
26 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type )
27 import TysPrim ( alphaTyVar )
28 import TysWiredIn ( nilDataCon, consDataCon, unitDataConId, tupleCon, mkListTy, mkTupleTy )
29 import Match ( matchSimply )
30 import PrelNames ( foldrName, buildName )
31 import SrcLoc ( noSrcLoc )
35 List comprehensions may be desugared in one of two ways: ``ordinary''
36 (as you would expect if you read SLPJ's book) and ``with foldr/build
37 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
39 There will be at least one ``qualifier'' in the input.
42 dsListComp :: [TypecheckedStmt]
43 -> Type -- Type of list elements
46 dsListComp quals elt_ty
47 | not opt_FoldrBuildOn -- Be boring
48 || isParallelComp quals
49 = deListComp quals (mkNilExpr elt_ty)
51 | otherwise -- foldr/build lives!
52 = newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
54 n_ty = mkTyVarTy n_tyvar
55 c_ty = mkFunTys [elt_ty, n_ty] n_ty
57 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
58 dfListComp c n quals `thenDs` \ result ->
59 dsLookupGlobalValue buildName `thenDs` \ build_id ->
60 returnDs (Var build_id `App` Type elt_ty
61 `App` mkLams [n_tyvar, c, n] result)
63 where isParallelComp (ParStmtOut bndrstmtss : _) = True
64 isParallelComp _ = False
67 %************************************************************************
69 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
71 %************************************************************************
73 Just as in Phil's chapter~7 in SLPJ, using the rules for
74 optimally-compiled list comprehensions. This is what Kevin followed
75 as well, and I quite happily do the same. The TQ translation scheme
76 transforms a list of qualifiers (either boolean expressions or
77 generators) into a single expression which implements the list
78 comprehension. Because we are generating 2nd-order polymorphic
79 lambda-calculus, calls to NIL and CONS must be applied to a type
80 argument, as well as their usual value arguments.
82 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
85 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
88 TQ << [ e | b , qs ] ++ L >> =
89 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
92 TQ << [ e | p <- L1, qs ] ++ L2 >> =
98 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
103 "h", "u1", "u2", and "u3" are new variables.
106 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
107 is the TE translation scheme. Note that we carry around the @L@ list
108 already desugared. @dsListComp@ does the top TE rule mentioned above.
110 To the above, we add an additional rule to deal with parallel list
111 comprehensions. The translation goes roughly as follows:
112 [ e | p1 <- e11, let v1 = e12, p2 <- e13
113 | q1 <- e21, let v2 = e22, q2 <- e23]
115 [ e | ((p1,v1,p2), (q1,v2,q2)) <-
116 zip [(p1,v1,p2) | p1 <- e11, let v1 = e12, p2 <- e13]
117 [(q1,v2,q2) | q1 <- e21, let v2 = e22, q2 <- e23]]
118 In the translation below, the ParStmtOut branch translates each parallel branch
119 into a sub-comprehension, and desugars each independently. The resulting lists
120 are fed to a zip function, we create a binding for all the variables bound in all
121 the comprehensions, and then we hand things off the the desugarer for bindings.
122 The zip function is generated here a) because it's small, and b) because then we
123 don't have to deal with arbitrary limits on the number of zip functions in the
124 prelude, nor which library the zip function came from.
125 The introduced tuples are Boxed, but only because I couldn't get it to work
126 with the Unboxed variety.
130 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
132 deListComp (ParStmtOut bndrstmtss : quals) list
133 = mapDs doListComp qualss `thenDs` \ exps ->
134 mapDs genAS bndrss `thenDs` \ ass ->
135 mapDs genA bndrss `thenDs` \ as ->
136 mapDs genAS' bndrss `thenDs` \ as's ->
137 let retTy = myTupleTy Boxed (length bndrss) qualTys
138 zipTy = foldr mkFunTy (mkListTy retTy) (map mkListTy qualTys)
140 newSysLocalDs zipTy `thenDs` \ zipFn ->
141 let target = mkConsExpr retTy (mkTupleExpr as) (foldl App (Var zipFn) (map Var as's))
142 zipExp = mkLet zipFn (zip4 bndrss ass as as's) exps target
144 deBindComp pat zipExp quals list
145 where (bndrss, stmtss) = unzip bndrstmtss
146 pats = map (\ps -> mkTuplePat (map VarPat ps)) bndrss
148 mkTuplePat ps = TuplePat ps Boxed
149 pat = TuplePat pats Boxed
151 qualss = map mkQuals bndrstmtss
152 mkQuals (bndrs, stmts) = (bndrs, stmts ++ [ExprStmt (myTupleExpr bndrs) noSrcLoc])
154 qualTys = map mkBndrsTy bndrss
155 mkBndrsTy bndrs = myTupleTy Boxed (length bndrs) (map idType bndrs)
157 doListComp (bndrs, stmts)
158 = dsListComp stmts (mkBndrsTy bndrs)
159 genA bndrs = newSysLocalDs (mkBndrsTy bndrs)
160 genAS bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
161 genAS' bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
163 mkLet zipFn vars exps target
165 foldr Lam (mkBody target vars) (map getAs vars))])
166 (foldl App (Var zipFn) exps)
167 getAs (_, as, _, _) = as
169 = foldr mkCase (foldr mkTuplCase target vars) vars
170 mkCase (ps, as, a, as') rest
171 = Case (Var as) as [(DataAlt nilDataCon, [], mkConApp nilDataCon []),
172 (DataAlt consDataCon, [a, as'], rest)]
173 mkTuplCase ([p], as, a, as') rest
174 = App (Lam p rest) (Var a)
175 mkTuplCase (ps, as, a, as') rest
176 = Case (Var a) a [(DataAlt (tupleCon Boxed (length ps)), ps, rest)]
178 myTupleTy boxity arity [ty] = ty
179 myTupleTy boxity arity tys = mkTupleTy boxity arity tys
180 myTupleExpr [] = HsVar unitDataConId
181 myTupleExpr [id] = HsVar id
182 myTupleExpr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
184 -- Last: the one to return
185 deListComp [ExprStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
186 = dsExpr expr `thenDs` \ core_expr ->
187 returnDs (mkConsExpr (exprType core_expr) core_expr list)
189 -- Non-last: must be a guard
190 deListComp (ExprStmt guard locn : quals) list -- rule B above
191 = dsExpr guard `thenDs` \ core_guard ->
192 deListComp quals list `thenDs` \ core_rest ->
193 returnDs (mkIfThenElse core_guard core_rest list)
195 -- [e | let B, qs] = let B in [e | qs]
196 deListComp (LetStmt binds : quals) list
197 = deListComp quals list `thenDs` \ core_rest ->
198 dsLet binds core_rest
200 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
201 = dsExpr list1 `thenDs` \ core_list1 ->
202 deBindComp pat core_list1 quals core_list2
204 deBindComp pat core_list1 quals core_list2
206 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
208 -- u1_ty is a [alpha] type, and u2_ty = alpha
209 u2_ty = outPatType pat
211 res_ty = exprType core_list2
212 h_ty = u1_ty `mkFunTy` res_ty
214 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
216 -- the "fail" value ...
218 core_fail = App (Var h) (Var u3)
219 letrec_body = App (Var h) core_list1
221 deListComp quals core_fail `thenDs` \ rest_expr ->
222 matchSimply (Var u2) ListComp pat
223 rest_expr core_fail `thenDs` \ core_match ->
226 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
227 (DataAlt consDataCon, [u2, u3], core_match)]
229 returnDs (Let (Rec [(h, rhs)]) letrec_body)
233 %************************************************************************
235 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
237 %************************************************************************
239 @dfListComp@ are the rules used with foldr/build turned on:
242 TE[ e | ] c n = c e n
243 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
244 TE[ e | p <- l , q ] c n = let
245 f = \ x b -> case x of
253 dfListComp :: Id -> Id -- 'c' and 'n'
254 -> [TypecheckedStmt] -- the rest of the qual's
257 -- Last: the one to return
258 dfListComp c_id n_id [ExprStmt expr locn]
259 = dsExpr expr `thenDs` \ core_expr ->
260 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
262 -- Non-last: must be a guard
263 dfListComp c_id n_id (ExprStmt guard locn : quals)
264 = dsExpr guard `thenDs` \ core_guard ->
265 dfListComp c_id n_id quals `thenDs` \ core_rest ->
266 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
268 dfListComp c_id n_id (LetStmt binds : quals)
269 -- new in 1.3, local bindings
270 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
271 dsLet binds core_rest
273 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
274 -- evaluate the two lists
275 = dsExpr list1 `thenDs` \ core_list1 ->
277 -- find the required type
278 let x_ty = outPatType pat
282 -- create some new local id's
283 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
285 -- build rest of the comprehesion
286 dfListComp c_id b quals `thenDs` \ core_rest ->
288 -- build the pattern match
289 matchSimply (Var x) ListComp pat core_rest (Var b) `thenDs` \ core_expr ->
291 -- now build the outermost foldr, and return
292 dsLookupGlobalValue foldrName `thenDs` \ foldr_id ->
294 Var foldr_id `App` Type x_ty
296 `App` mkLams [x, b] core_expr