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, outPatType )
18 import DsMonad -- the monadery used in the desugarer
21 import CmdLineOpts ( opt_FoldrBuildOn )
22 import CoreUtils ( exprType, mkIfThenElse )
25 import Type ( mkTyVarTy, mkFunTys, mkFunTy, Type )
26 import TysPrim ( alphaTyVar )
27 import TysWiredIn ( nilDataCon, consDataCon, unitDataConId, mkListTy, mkTupleTy )
28 import Match ( matchSimply )
29 import PrelNames ( foldrName, buildName )
30 import SrcLoc ( noSrcLoc )
33 List comprehensions may be desugared in one of two ways: ``ordinary''
34 (as you would expect if you read SLPJ's book) and ``with foldr/build
35 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
37 There will be at least one ``qualifier'' in the input.
40 dsListComp :: [TypecheckedStmt]
41 -> Type -- Type of list elements
44 dsListComp quals elt_ty
45 | not opt_FoldrBuildOn -- Be boring
46 || isParallelComp quals
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] ->
56 dfListComp c n quals `thenDs` \ result ->
57 dsLookupGlobalValue buildName `thenDs` \ build_id ->
58 returnDs (Var build_id `App` Type elt_ty
59 `App` mkLams [n_tyvar, c, n] result)
61 where isParallelComp (ParStmtOut bndrstmtss : _) = True
62 isParallelComp _ = False
65 %************************************************************************
67 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
69 %************************************************************************
71 Just as in Phil's chapter~7 in SLPJ, using the rules for
72 optimally-compiled list comprehensions. This is what Kevin followed
73 as well, and I quite happily do the same. The TQ translation scheme
74 transforms a list of qualifiers (either boolean expressions or
75 generators) into a single expression which implements the list
76 comprehension. Because we are generating 2nd-order polymorphic
77 lambda-calculus, calls to NIL and CONS must be applied to a type
78 argument, as well as their usual value arguments.
80 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
83 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
86 TQ << [ e | b , qs ] ++ L >> =
87 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
90 TQ << [ e | p <- L1, qs ] ++ L2 >> =
96 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
101 "h", "u1", "u2", and "u3" are new variables.
104 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
105 is the TE translation scheme. Note that we carry around the @L@ list
106 already desugared. @dsListComp@ does the top TE rule mentioned above.
108 To the above, we add an additional rule to deal with parallel list
109 comprehensions. The translation goes roughly as follows:
110 [ e | p1 <- e11, let v1 = e12, p2 <- e13
111 | q1 <- e21, let v2 = e22, q2 <- e23]
113 [ e | ((x1, .., xn), (y1, ..., ym)) <-
114 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
115 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
116 where (x1, .., xn) are the variables bound in p1, v1, p2
117 (y1, .., ym) are the variables bound in q1, v2, q2
119 In the translation below, the ParStmtOut branch translates each parallel branch
120 into a sub-comprehension, and desugars each independently. The resulting lists
121 are fed to a zip function, we create a binding for all the variables bound in all
122 the comprehensions, and then we hand things off the the desugarer for bindings.
123 The zip function is generated here a) because it's small, and b) because then we
124 don't have to deal with arbitrary limits on the number of zip functions in the
125 prelude, nor which library the zip function came from.
126 The introduced tuples are Boxed, but only because I couldn't get it to work
127 with the Unboxed variety.
131 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
133 deListComp (ParStmtOut bndrstmtss : quals) list
134 = mapDs do_list_comp bndrstmtss `thenDs` \ exps ->
135 mkZipBind qual_tys `thenDs` \ (zip_fn, zip_rhs) ->
137 -- Deal with [e | pat <- zip l1 .. ln] in example above
138 deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
141 where -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
142 pat = TuplePat pats Boxed
143 pats = map (\(bs,_) -> mk_hs_tuple_pat bs) bndrstmtss
145 -- Types of (x1,..,xn), (y1,..,yn) etc
146 qual_tys = [ mk_bndrs_tys bndrs | (bndrs,_) <- bndrstmtss ]
148 do_list_comp (bndrs, stmts)
149 = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
152 mk_bndrs_tys bndrs = mk_tuple_ty (map idType bndrs)
154 -- Last: the one to return
155 deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
156 = dsExpr expr `thenDs` \ core_expr ->
157 returnDs (mkConsExpr (exprType core_expr) core_expr list)
159 -- Non-last: must be a guard
160 deListComp (ExprStmt guard ty locn : quals) list -- rule B above
161 = dsExpr guard `thenDs` \ core_guard ->
162 deListComp quals list `thenDs` \ core_rest ->
163 returnDs (mkIfThenElse core_guard core_rest list)
165 -- [e | let B, qs] = let B in [e | qs]
166 deListComp (LetStmt binds : quals) list
167 = deListComp quals list `thenDs` \ core_rest ->
168 dsLet binds core_rest
170 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
171 = dsExpr list1 `thenDs` \ core_list1 ->
172 deBindComp pat core_list1 quals core_list2
177 deBindComp pat core_list1 quals core_list2
179 u3_ty@u1_ty = exprType core_list1 -- two names, same thing
181 -- u1_ty is a [alpha] type, and u2_ty = alpha
182 u2_ty = outPatType pat
184 res_ty = exprType core_list2
185 h_ty = u1_ty `mkFunTy` res_ty
187 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
189 -- the "fail" value ...
191 core_fail = App (Var h) (Var u3)
192 letrec_body = App (Var h) core_list1
194 deListComp quals core_fail `thenDs` \ rest_expr ->
195 matchSimply (Var u2) (DoCtxt ListComp) pat
196 rest_expr core_fail `thenDs` \ core_match ->
199 Case (Var u1) u1 [(DataAlt nilDataCon, [], core_list2),
200 (DataAlt consDataCon, [u2, u3], core_match)]
202 returnDs (Let (Rec [(h, rhs)]) letrec_body)
207 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
208 -- mkZipBind [t1, t2]
209 -- = (zip, \as1:[t1] as2:[t2]
212 -- (a1:as'1) -> case as2 of
214 -- (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
217 = mapDs newSysLocalDs list_tys `thenDs` \ ass ->
218 mapDs newSysLocalDs elt_tys `thenDs` \ as' ->
219 mapDs newSysLocalDs list_tys `thenDs` \ as's ->
220 newSysLocalDs zip_fn_ty `thenDs` \ zip_fn ->
222 inner_rhs = mkConsExpr ret_elt_ty (mkTupleExpr as') (mkVarApps (Var zip_fn) as's)
223 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
225 returnDs (zip_fn, mkLams ass zip_body)
227 list_tys = map mkListTy elt_tys
228 ret_elt_ty = mk_tuple_ty elt_tys
229 zip_fn_ty = mkFunTys list_tys (mkListTy ret_elt_ty)
231 mk_case (as, a', as') rest
232 = Case (Var as) as [(DataAlt nilDataCon, [], mkNilExpr ret_elt_ty),
233 (DataAlt consDataCon, [a', as'], rest)]
236 mk_tuple_ty :: [Type] -> Type
237 mk_tuple_ty [ty] = ty
238 mk_tuple_ty tys = mkTupleTy Boxed (length tys) tys
240 -- Helper functions that makes an HsTuple only for non-1-sized tuples
241 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
242 mk_hs_tuple_expr [] = HsVar unitDataConId
243 mk_hs_tuple_expr [id] = HsVar id
244 mk_hs_tuple_expr ids = ExplicitTuple [ HsVar i | i <- ids ] Boxed
246 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
247 mk_hs_tuple_pat [b] = VarPat b
248 mk_hs_tuple_pat bs = TuplePat (map VarPat bs) Boxed
252 %************************************************************************
254 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
256 %************************************************************************
258 @dfListComp@ are the rules used with foldr/build turned on:
261 TE[ e | ] c n = c e n
262 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
263 TE[ e | p <- l , q ] c n = let
264 f = \ x b -> case x of
272 dfListComp :: Id -> Id -- 'c' and 'n'
273 -> [TypecheckedStmt] -- the rest of the qual's
276 -- Last: the one to return
277 dfListComp c_id n_id [ResultStmt expr locn]
278 = dsExpr expr `thenDs` \ core_expr ->
279 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
281 -- Non-last: must be a guard
282 dfListComp c_id n_id (ExprStmt guard ty locn : quals)
283 = dsExpr guard `thenDs` \ core_guard ->
284 dfListComp c_id n_id quals `thenDs` \ core_rest ->
285 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
287 dfListComp c_id n_id (LetStmt binds : quals)
288 -- new in 1.3, local bindings
289 = dfListComp c_id n_id quals `thenDs` \ core_rest ->
290 dsLet binds core_rest
292 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
293 -- evaluate the two lists
294 = dsExpr list1 `thenDs` \ core_list1 ->
296 -- find the required type
297 let x_ty = outPatType pat
301 -- create some new local id's
302 newSysLocalsDs [b_ty,x_ty] `thenDs` \ [b,x] ->
304 -- build rest of the comprehesion
305 dfListComp c_id b quals `thenDs` \ core_rest ->
307 -- build the pattern match
308 matchSimply (Var x) (DoCtxt ListComp)
309 pat core_rest (Var b) `thenDs` \ core_expr ->
311 -- now build the outermost foldr, and return
312 dsLookupGlobalValue foldrName `thenDs` \ foldr_id ->
314 Var foldr_id `App` Type x_ty
316 `App` mkLams [x, b] core_expr