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(..), HsDoContext(..) )
15 import TcHsSyn ( TypecheckedStmt, TypecheckedPat, TypecheckedHsExpr )
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, mkListTy, mkTupleTy )
29 import Match ( matchSimply )
30 import PrelNames ( foldrName, buildName )
31 import SrcLoc ( noSrcLoc )
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 || isParallelComp quals
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)
62 where isParallelComp (ParStmtOut bndrstmtss : _) = True
63 isParallelComp _ = False
66 %************************************************************************
68 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
70 %************************************************************************
72 Just as in Phil's chapter~7 in SLPJ, using the rules for
73 optimally-compiled list comprehensions. This is what Kevin followed
74 as well, and I quite happily do the same. The TQ translation scheme
75 transforms a list of qualifiers (either boolean expressions or
76 generators) into a single expression which implements the list
77 comprehension. Because we are generating 2nd-order polymorphic
78 lambda-calculus, calls to NIL and CONS must be applied to a type
79 argument, as well as their usual value arguments.
81 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
84 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
87 TQ << [ e | b , qs ] ++ L >> =
88 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
91 TQ << [ e | p <- L1, qs ] ++ L2 >> =
97 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
102 "h", "u1", "u2", and "u3" are new variables.
105 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
106 is the TE translation scheme. Note that we carry around the @L@ list
107 already desugared. @dsListComp@ does the top TE rule mentioned above.
109 To the above, we add an additional rule to deal with parallel list
110 comprehensions. The translation goes roughly as follows:
111 [ e | p1 <- e11, let v1 = e12, p2 <- e13
112 | q1 <- e21, let v2 = e22, q2 <- e23]
114 [ e | ((x1, .., xn), (y1, ..., ym)) <-
115 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
116 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
117 where (x1, .., xn) are the variables bound in p1, v1, p2
118 (y1, .., ym) are the variables bound in q1, v2, q2
120 In the translation below, the ParStmtOut branch translates each parallel branch
121 into a sub-comprehension, and desugars each independently. The resulting lists
122 are fed to a zip function, we create a binding for all the variables bound in all
123 the comprehensions, and then we hand things off the the desugarer for bindings.
124 The zip function is generated here a) because it's small, and b) because then we
125 don't have to deal with arbitrary limits on the number of zip functions in the
126 prelude, nor which library the zip function came from.
127 The introduced tuples are Boxed, but only because I couldn't get it to work
128 with the Unboxed variety.
132 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
134 deListComp (ParStmtOut bndrstmtss : quals) list
135 = mapDs do_list_comp bndrstmtss `thenDs` \ exps ->
136 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
138 -- Deal with [e | pat <- zip l1 .. ln] in example above
139 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
142 where -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
143 pat = TuplePat pats Boxed
144 pats = map (\(bs,_) -> mk_hs_tuple_pat bs) bndrstmtss
146 -- Types of (x1,..,xn), (y1,..,yn) etc
147 qual_tys = [ mk_bndrs_tys bndrs | (bndrs,_) <- bndrstmtss ]
149 do_list_comp (bndrs, stmts)
150 = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
153 mk_bndrs_tys bndrs = mk_tuple_ty (map idType bndrs)
155 -- Last: the one to return
156 deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
157 = dsExpr expr `thenDs` \ core_expr ->
158 returnDs (mkConsExpr (exprType core_expr) core_expr list)
160 -- Non-last: must be a guard
161 deListComp (ExprStmt guard locn : quals) list -- rule B above
162 = dsExpr guard `thenDs` \ core_guard ->
163 deListComp quals list `thenDs` \ core_rest ->
164 returnDs (mkIfThenElse core_guard core_rest list)
166 -- [e | let B, qs] = let B in [e | qs]
167 deListComp (LetStmt binds : quals) list
168 = deListComp quals list `thenDs` \ core_rest ->
169 dsLet binds core_rest
171 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
172 = dsExpr list1 `thenDs` \ core_list1 ->
173 deBindComp pat core_list1 quals core_list2
178 deBindComp pat core_list1 quals core_list2
180 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
182 -- u1_ty is a [alpha] type, and u2_ty = alpha
183 u2_ty = outPatType pat
185 res_ty = exprType core_list2
186 h_ty = u1_ty `mkFunTy` res_ty
188 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
190 -- the "fail" value ...
192 core_fail = App (Var h) (Var u3)
193 letrec_body = App (Var h) core_list1
195 deListComp quals core_fail `thenDs` \ rest_expr ->
196 matchSimply (Var u2) (DoCtxt ListComp) pat
197 rest_expr core_fail `thenDs` \ core_match ->
200 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
201 (DataAlt consDataCon, [u2, u3], core_match)]
203 returnDs (Let (Rec [(h, rhs)]) letrec_body)
208 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
209 -- mkZipBind [t1, t2]
210 -- = (zip, \as1:[t1] as2:[t2]
213 -- (a1:as'1) -> case as2 of
215 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
218 = mapDs newSysLocalDs list_tys `thenDs` \ ass ->
219 mapDs newSysLocalDs elt_tys `thenDs` \ as' ->
220 mapDs newSysLocalDs list_tys `thenDs` \ as's ->
221 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
223 inner_rhs = mkConsExpr ret_elt_ty (mkTupleExpr as') (mkVarApps (Var zip_fn) as's)
224 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
226 returnDs (zip_fn, mkLams ass zip_body)
228 list_tys = map mkListTy elt_tys
229 ret_elt_ty = mk_tuple_ty elt_tys
230 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
232 mk_case (as, a', as') rest
233 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
234 (DataAlt consDataCon, [a', as'], rest)]
237 mk_tuple_ty :: [Type] -> Type
238 mk_tuple_ty [ty] = ty
239 mk_tuple_ty tys = mkTupleTy Boxed (length tys) tys
241 -- Helper functions that makes an HsTuple only for non-1-sized tuples
242 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
243 mk_hs_tuple_expr [] = HsVar unitDataConId
244 mk_hs_tuple_expr [id] = HsVar id
245 mk_hs_tuple_expr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
247 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
248 mk_hs_tuple_pat [b] = VarPat b
249 mk_hs_tuple_pat bs = TuplePat (map VarPat bs) Boxed
253 %************************************************************************
255 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
257 %************************************************************************
259 @dfListComp@ are the rules used with foldr/build turned on:
262 TE[ e | ] c n = c e n
263 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
264 TE[ e | p <- l , q ] c n = let
265 f = \ x b -> case x of
273 dfListComp :: Id -> Id -- 'c' and 'n'
274 -> [TypecheckedStmt] -- the rest of the qual's
277 -- Last: the one to return
278 dfListComp c_id n_id [ResultStmt expr locn]
279 = dsExpr expr `thenDs` \ core_expr ->
280 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
282 -- Non-last: must be a guard
283 dfListComp c_id n_id (ExprStmt guard locn : quals)
284 = dsExpr guard `thenDs` \ core_guard ->
285 dfListComp c_id n_id quals `thenDs` \ core_rest ->
286 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
288 dfListComp c_id n_id (LetStmt binds : quals)
289 -- new in 1.3, local bindings
290 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
291 dsLet binds core_rest
293 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
294 -- evaluate the two lists
295 = dsExpr list1 `thenDs` \ core_list1 ->
297 -- find the required type
298 let x_ty = outPatType pat
302 -- create some new local id's
303 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
305 -- build rest of the comprehesion
306 dfListComp c_id b quals `thenDs` \ core_rest ->
308 -- build the pattern match
309 matchSimply (Var x) (DoCtxt ListComp)
310 pat core_rest (Var b) `thenDs` \ core_expr ->
312 -- now build the outermost foldr, and return
313 dsLookupGlobalValue foldrName `thenDs` \ foldr_id ->
315 Var foldr_id `App` Type x_ty
317 `App` mkLams [x, b] core_expr