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 = deListComp quals (mkNilExpr elt_ty)
50 | otherwise -- foldr/build lives!
51 = newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
53 n_ty = mkTyVarTy n_tyvar
54 c_ty = mkFunTys [elt_ty, n_ty] n_ty
56 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
57 dfListComp c n quals `thenDs` \ result ->
58 dsLookupGlobalValue buildName `thenDs` \ build_id ->
59 returnDs (Var build_id `App` Type elt_ty
60 `App` mkLams [n_tyvar, c, n] result)
63 %************************************************************************
65 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
67 %************************************************************************
69 Just as in Phil's chapter~7 in SLPJ, using the rules for
70 optimally-compiled list comprehensions. This is what Kevin followed
71 as well, and I quite happily do the same. The TQ translation scheme
72 transforms a list of qualifiers (either boolean expressions or
73 generators) into a single expression which implements the list
74 comprehension. Because we are generating 2nd-order polymorphic
75 lambda-calculus, calls to NIL and CONS must be applied to a type
76 argument, as well as their usual value arguments.
78 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
81 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
84 TQ << [ e | b , qs ] ++ L >> =
85 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
88 TQ << [ e | p <- L1, qs ] ++ L2 >> =
94 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
99 "h", "u1", "u2", and "u3" are new variables.
102 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
103 is the TE translation scheme. Note that we carry around the @L@ list
104 already desugared. @dsListComp@ does the top TE rule mentioned above.
106 To the above, we add an additional rule to deal with parallel list
107 comprehensions. The translation goes roughly as follows:
108 [ e | p1 <- e11, let v1 = e12, p2 <- e13
109 | q1 <- e21, let v2 = e22, q2 <- e23]
111 [ e | ((p1,v1,p2), (q1,v2,q2)) <-
112 zip [(p1,v1,p2) | p1 <- e11, let v1 = e12, p2 <- e13]
113 [(q1,v2,q2) | q1 <- e21, let v2 = e22, q2 <- e23]]
114 In the translation below, the ParStmtOut branch translates each parallel branch
115 into a sub-comprehension, and desugars each independently. The resulting lists
116 are fed to a zip function, we create a binding for all the variables bound in all
117 the comprehensions, and then we hand things off the the desugarer for bindings.
118 The zip function is generated here a) because it's small, and b) because then we
119 don't have to deal with arbitrary limits on the number of zip functions in the
120 prelude, nor which library the zip function came from.
121 The introduced tuples are Boxed, but only because I couldn't get it to work
122 with the Unboxed variety.
126 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
128 deListComp (ParStmtOut bndrstmtss : quals) list
129 = mapDs doListComp qualss `thenDs` \ exps ->
130 mapDs genAS bndrss `thenDs` \ ass ->
131 mapDs genA bndrss `thenDs` \ as ->
132 mapDs genAS' bndrss `thenDs` \ as's ->
133 let retTy = myTupleTy Boxed (length bndrss) qualTys
134 zipTy = foldr mkFunTy (mkListTy retTy) (map mkListTy qualTys)
136 newSysLocalDs zipTy `thenDs` \ zipFn ->
137 let target = mkConsExpr retTy (mkTupleExpr as) (foldl App (Var zipFn) (map Var as's))
138 zipExp = mkLet zipFn (zip4 bndrss ass as as's) exps target
140 deBindComp pat zipExp quals list
141 where (bndrss, stmtss) = unzip bndrstmtss
142 pats = map (\ps -> mkTuplePat (map VarPat ps)) bndrss
144 mkTuplePat ps = TuplePat ps Boxed
145 pat = TuplePat pats Boxed
147 qualss = map mkQuals bndrstmtss
148 mkQuals (bndrs, stmts) = (bndrs, stmts ++ [ExprStmt (myTupleExpr bndrs) noSrcLoc])
150 qualTys = map mkBndrsTy bndrss
151 mkBndrsTy bndrs = myTupleTy Boxed (length bndrs) (map idType bndrs)
153 doListComp (bndrs, stmts)
154 = dsListComp stmts (mkBndrsTy bndrs)
155 genA bndrs = newSysLocalDs (mkBndrsTy bndrs)
156 genAS bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
157 genAS' bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
159 mkLet zipFn vars exps target
161 foldr Lam (mkBody target vars) (map getAs vars))])
162 (foldl App (Var zipFn) exps)
163 getAs (_, as, _, _) = as
165 = foldr mkCase (foldr mkTuplCase target vars) vars
166 mkCase (ps, as, a, as') rest
167 = Case (Var as) as [(DataAlt nilDataCon, [], mkConApp nilDataCon []),
168 (DataAlt consDataCon, [a, as'], rest)]
169 mkTuplCase ([p], as, a, as') rest
170 = App (Lam p rest) (Var a)
171 mkTuplCase (ps, as, a, as') rest
172 = Case (Var a) a [(DataAlt (tupleCon Boxed (length ps)), ps, rest)]
174 myTupleTy boxity arity [ty] = ty
175 myTupleTy boxity arity tys = mkTupleTy boxity arity tys
176 myTupleExpr [] = HsVar unitDataConId
177 myTupleExpr [id] = HsVar id
178 myTupleExpr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
180 -- Last: the one to return
181 deListComp [ExprStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
182 = dsExpr expr `thenDs` \ core_expr ->
183 returnDs (mkConsExpr (exprType core_expr) core_expr list)
185 -- Non-last: must be a guard
186 deListComp (ExprStmt guard locn : quals) list -- rule B above
187 = dsExpr guard `thenDs` \ core_guard ->
188 deListComp quals list `thenDs` \ core_rest ->
189 returnDs (mkIfThenElse core_guard core_rest list)
191 -- [e | let B, qs] = let B in [e | qs]
192 deListComp (LetStmt binds : quals) list
193 = deListComp quals list `thenDs` \ core_rest ->
194 dsLet binds core_rest
196 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
197 = dsExpr list1 `thenDs` \ core_list1 ->
198 deBindComp pat core_list1 quals core_list2
200 deBindComp pat core_list1 quals core_list2
202 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
204 -- u1_ty is a [alpha] type, and u2_ty = alpha
205 u2_ty = outPatType pat
207 res_ty = exprType core_list2
208 h_ty = u1_ty `mkFunTy` res_ty
210 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
212 -- the "fail" value ...
214 core_fail = App (Var h) (Var u3)
215 letrec_body = App (Var h) core_list1
217 deListComp quals core_fail `thenDs` \ rest_expr ->
218 matchSimply (Var u2) ListComp pat
219 rest_expr core_fail `thenDs` \ core_match ->
222 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
223 (DataAlt consDataCon, [u2, u3], core_match)]
225 returnDs (Let (Rec [(h, rhs)]) letrec_body)
229 %************************************************************************
231 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
233 %************************************************************************
235 @dfListComp@ are the rules used with foldr/build turned on:
238 TE[ e | ] c n = c e n
239 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
240 TE[ e | p <- l , q ] c n = let
241 f = \ x b -> case x of
249 dfListComp :: Id -> Id -- 'c' and 'n'
250 -> [TypecheckedStmt] -- the rest of the qual's
253 -- Last: the one to return
254 dfListComp c_id n_id [ExprStmt expr locn]
255 = dsExpr expr `thenDs` \ core_expr ->
256 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
258 -- Non-last: must be a guard
259 dfListComp c_id n_id (ExprStmt guard locn : quals)
260 = dsExpr guard `thenDs` \ core_guard ->
261 dfListComp c_id n_id quals `thenDs` \ core_rest ->
262 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
264 dfListComp c_id n_id (LetStmt binds : quals)
265 -- new in 1.3, local bindings
266 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
267 dsLet binds core_rest
269 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
270 -- evaluate the two lists
271 = dsExpr list1 `thenDs` \ core_list1 ->
273 -- find the required type
274 let x_ty = outPatType pat
278 -- create some new local id's
279 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
281 -- build rest of the comprehesion
282 dfListComp c_id b quals `thenDs` \ core_rest ->
284 -- build the pattern match
285 matchSimply (Var x) ListComp pat core_rest (Var b) `thenDs` \ core_expr ->
287 -- now build the outermost foldr, and return
288 dsLookupGlobalValue foldrName `thenDs` \ foldr_id ->
290 Var foldr_id `App` Type x_ty
292 `App` mkLams [x, b] core_expr