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(..) )
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 )
34 List comprehensions may be desugared in one of two ways: ``ordinary''
35 (as you would expect if you read SLPJ's book) and ``with foldr/build
36 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
38 There will be at least one ``qualifier'' in the input.
41 dsListComp :: [TypecheckedStmt]
42 -> Type -- Type of list elements
45 dsListComp quals elt_ty
46 | not opt_FoldrBuildOn -- Be boring
47 = deListComp quals (mkNilExpr elt_ty)
49 | otherwise -- foldr/build lives!
50 = newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
52 n_ty = mkTyVarTy n_tyvar
53 c_ty = mkFunTys [elt_ty, n_ty] n_ty
55 newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
57 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)
64 %************************************************************************
66 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
68 %************************************************************************
70 Just as in Phil's chapter~7 in SLPJ, using the rules for
71 optimally-compiled list comprehensions. This is what Kevin followed
72 as well, and I quite happily do the same. The TQ translation scheme
73 transforms a list of qualifiers (either boolean expressions or
74 generators) into a single expression which implements the list
75 comprehension. Because we are generating 2nd-order polymorphic
76 lambda-calculus, calls to NIL and CONS must be applied to a type
77 argument, as well as their usual value arguments.
79 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
82 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
85 TQ << [ e | b , qs ] ++ L >> =
86 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
89 TQ << [ e | p <- L1, qs ] ++ L2 >> =
95 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
100 "h", "u1", "u2", and "u3" are new variables.
103 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
104 is the TE translation scheme. Note that we carry around the @L@ list
105 already desugared. @dsListComp@ does the top TE rule mentioned above.
107 To the above, we add an additional rule to deal with parallel list
108 comprehensions. The translation goes roughly as follows:
109 [ e | p1 <- e11, let v1 = e12, p2 <- e13
110 | q1 <- e21, let v2 = e22, q2 <- e23]
112 [ e | ((p1,v1,p2), (q1,v2,q2)) <-
113 zip [(p1,v1,p2) | p1 <- e11, let v1 = e12, p2 <- e13]
114 [(q1,v2,q2) | q1 <- e21, let v2 = e22, q2 <- e23]]
115 In the translation below, the ParStmtOut branch translates each parallel branch
116 into a sub-comprehension, and desugars each independently. The resulting lists
117 are fed to a zip function, we create a binding for all the variables bound in all
118 the comprehensions, and then we hand things off the the desugarer for bindings.
119 The zip function is generated here a) because it's small, and b) because then we
120 don't have to deal with arbitrary limits on the number of zip functions in the
121 prelude, nor which library the zip function came from.
122 The introduced tuples are Boxed, but only because I couldn't get it to work
123 with the Unboxed variety.
127 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
129 deListComp (ParStmtOut bndrstmtss : quals) list
130 = mapDs doListComp qualss `thenDs` \ exps ->
131 mapDs genAS bndrss `thenDs` \ ass ->
132 mapDs genA bndrss `thenDs` \ as ->
133 mapDs genAS' bndrss `thenDs` \ as's ->
134 let retTy = myTupleTy Boxed (length bndrss) qualTys
135 zipTy = foldr mkFunTy (mkListTy retTy) (map mkListTy qualTys)
137 newSysLocalDs zipTy `thenDs` \ zipFn ->
138 let target = mkConsExpr retTy (mkTupleExpr as) (foldl App (Var zipFn) (map Var as's))
139 zipExp = mkLet zipFn (zip4 bndrss ass as as's) exps target
141 deBindComp pat zipExp quals list
142 where (bndrss, stmtss) = unzip bndrstmtss
143 pats = map (\ps -> mkTuplePat (map VarPat ps)) bndrss
145 mkTuplePat ps = TuplePat ps Boxed
146 pat = TuplePat pats Boxed
148 qualss = map mkQuals bndrstmtss
149 mkQuals (bndrs, stmts) = (bndrs, stmts ++ [ReturnStmt (myTupleExpr bndrs)])
151 qualTys = map mkBndrsTy bndrss
152 mkBndrsTy bndrs = myTupleTy Boxed (length bndrs) (map idType bndrs)
154 doListComp (bndrs, stmts)
155 = dsListComp stmts (mkBndrsTy bndrs)
156 genA bndrs = newSysLocalDs (mkBndrsTy bndrs)
157 genAS bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
158 genAS' bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
160 mkLet zipFn vars exps target
162 foldr Lam (mkBody target vars) (map getAs vars))])
163 (foldl App (Var zipFn) exps)
164 getAs (_, as, _, _) = as
166 = foldr mkCase (foldr mkTuplCase target vars) vars
167 mkCase (ps, as, a, as') rest
168 = Case (Var as) as [(DataAlt nilDataCon, [], mkConApp nilDataCon []),
169 (DataAlt consDataCon, [a, as'], rest)]
170 mkTuplCase ([p], as, a, as') rest
171 = App (Lam p rest) (Var a)
172 mkTuplCase (ps, as, a, as') rest
173 = Case (Var a) a [(DataAlt (tupleCon Boxed (length ps)), ps, rest)]
175 myTupleTy boxity arity [ty] = ty
176 myTupleTy boxity arity tys = mkTupleTy boxity arity tys
177 myTupleExpr [] = HsVar unitDataConId
178 myTupleExpr [id] = HsVar id
179 myTupleExpr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
181 deListComp [ReturnStmt expr] 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 deListComp (GuardStmt guard locn : quals) list -- rule B above
186 = dsExpr guard `thenDs` \ core_guard ->
187 deListComp quals list `thenDs` \ core_rest ->
188 returnDs (mkIfThenElse core_guard core_rest list)
190 -- [e | let B, qs] = let B in [e | qs]
191 deListComp (LetStmt binds : quals) list
192 = deListComp quals list `thenDs` \ core_rest ->
193 dsLet binds core_rest
195 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
196 = dsExpr list1 `thenDs` \ core_list1 ->
197 deBindComp pat core_list1 quals core_list2
199 deBindComp pat core_list1 quals core_list2
201 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
203 -- u1_ty is a [alpha] type, and u2_ty = alpha
204 u2_ty = outPatType pat
206 res_ty = exprType core_list2
207 h_ty = u1_ty `mkFunTy` res_ty
209 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
211 -- the "fail" value ...
213 core_fail = App (Var h) (Var u3)
214 letrec_body = App (Var h) core_list1
216 deListComp quals core_fail `thenDs` \ rest_expr ->
217 matchSimply (Var u2) ListCompMatch pat
218 rest_expr core_fail `thenDs` \ core_match ->
221 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
222 (DataAlt consDataCon, [u2, u3], core_match)]
224 returnDs (Let (Rec [(h, rhs)]) letrec_body)
228 %************************************************************************
230 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
232 %************************************************************************
234 @dfListComp@ are the rules used with foldr/build turned on:
237 TE[ e | ] c n = c e n
238 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
239 TE[ e | p <- l , q ] c n = let
240 f = \ x b -> case x of
248 dfListComp :: Id -> Id -- 'c' and 'n'
249 -> [TypecheckedStmt] -- the rest of the qual's
252 dfListComp c_id n_id [ReturnStmt expr]
253 = dsExpr expr `thenDs` \ core_expr ->
254 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
256 dfListComp c_id n_id (GuardStmt guard locn : quals)
257 = dsExpr guard `thenDs` \ core_guard ->
258 dfListComp c_id n_id quals `thenDs` \ core_rest ->
259 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
261 dfListComp c_id n_id (LetStmt binds : quals)
262 -- new in 1.3, local bindings
263 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
264 dsLet binds core_rest
266 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
267 -- evaluate the two lists
268 = dsExpr list1 `thenDs` \ core_list1 ->
270 -- find the required type
271 let x_ty = outPatType pat
275 -- create some new local id's
276 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
278 -- build rest of the comprehesion
279 dfListComp c_id b quals `thenDs` \ core_rest ->
281 -- build the pattern match
282 matchSimply (Var x) ListCompMatch pat core_rest (Var b) `thenDs` \ core_expr ->
284 -- now build the outermost foldr, and return
285 dsLookupGlobalValue foldrName `thenDs` \ foldr_id ->
287 Var foldr_id `App` Type x_ty
289 `App` mkLams [x, b] core_expr