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 HsSyn ( Stmt(..), HsExpr )
14 import TcHsSyn ( TypecheckedStmt, TypecheckedHsExpr )
15 import DsHsSyn ( outPatType )
18 import DsMonad -- the monadery used in the desugarer
21 import CmdLineOpts ( opt_FoldrBuildOn )
22 import CoreUtils ( coreExprType )
23 import Var ( Id, TyVar )
24 import Const ( Con(..) )
25 import PrelInfo ( foldrId )
26 import Type ( mkTyVarTy, mkForAllTy, mkFunTys, mkFunTy, Type )
27 import TysPrim ( alphaTyVar, alphaTy )
28 import TysWiredIn ( nilDataCon, consDataCon, listTyCon )
29 import Match ( matchSimply )
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 = deListComp quals nil_expr
48 | otherwise -- foldr/build lives!
49 = newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
51 alpha_to_alpha = alphaTy `mkFunTy` alphaTy
53 n_ty = mkTyVarTy n_tyvar
54 c_ty = mkFunTys [elt_ty, n_ty] n_ty
55 g_ty = mkForAllTy alphaTyVar (
56 (elt_ty `mkFunTy` alpha_to_alpha)
61 newSysLocalsDs [c_ty,n_ty,g_ty] `thenDs` \ [c, n, g] ->
65 quals `thenDs` \ result ->
67 returnDs (mkBuild elt_ty n_tyvar c n g result)
69 nil_expr = mkNilExpr elt_ty
72 %************************************************************************
74 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
76 %************************************************************************
78 Just as in Phil's chapter~7 in SLPJ, using the rules for
79 optimally-compiled list comprehensions. This is what Kevin followed
80 as well, and I quite happily do the same. The TQ translation scheme
81 transforms a list of qualifiers (either boolean expressions or
82 generators) into a single expression which implements the list
83 comprehension. Because we are generating 2nd-order polymorphic
84 lambda-calculus, calls to NIL and CONS must be applied to a type
85 argument, as well as their usual value arguments.
87 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
90 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
93 TQ << [ e | b , qs ] ++ L >> =
94 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
97 TQ << [ e | p <- L1, qs ] ++ L2 >> =
103 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
108 "h", "u1", "u2", and "u3" are new variables.
111 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
112 is the TE translation scheme. Note that we carry around the @L@ list
113 already desugared. @dsListComp@ does the top TE rule mentioned above.
115 deListComp :: [TypecheckedStmt]
116 -> CoreExpr -> CoreExpr -- Cons and nil resp; can be copied freely
119 deListComp [ReturnStmt expr] cons nil
120 = dsExpr expr `thenDs` \ expr' ->
121 returnDs (mkApps cons [expr', nil])
123 deListComp (GuardStmt guard locn : quals) cons nil
124 = dsExpr guard `thenDs` \ guard' ->
125 deListComp quals cons nil `thenDs` \ rest' ->
126 returnDs (mkIfThenElse guard' rest' nil)
128 deListComp (LetStmt binds : quals) cons nil
129 = deListComp quals cons nil `thenDs` \ rest' ->
132 deListComp (BindStmt pat list locn : quals) cons nil
133 = dsExpr list `thenDs` \ list' ->
135 pat_ty = outPatType pat
136 nil_ty = coreExprType nil
138 newSysLocalsDs [pat_ty, nil_ty] `thenDs` \ [x,ys] ->
140 dsListComp quals cons (Var ys) `thenDs` \ rest ->
141 matchSimply (Var x) ListCompMatch pat
142 rest (Var ys) `thenDs` \ core_match ->
143 bindNonRecDs (mkLams [x,ys] fn_body) $ \ fn ->
144 dsListExpr list (Var fn) nil
147 data FExpr = FEOther CoreExpr -- Default case
149 | FEConsComposedWith CoreExpr -- (cons . e)
152 feComposeWith FECons g
153 = returnDs (FEConsComposedWith g)
155 feComposeWith (FEOther f) g
156 = composeWith f f `thenDs` \ h ->
159 feComposeWith (FEConsComposedWith f) g
160 = composeWith f f `thenDs` \ h ->
161 returnDs (FEConsComposedWith h)
165 = newSysLocalDs arg_ty `thenDs` \ x ->
166 returnDs (Lam x (App e (App f (Var x))))
168 arg_ty = case splitFunTy_maybe (coreExprType g) of
169 Just (arg_ty,_) -> arg_ty
170 other -> panic "feComposeWith"
172 deListExpr :: TypecheckedHsExpr
173 -> FExpr -> FExpr -- Cons and nil expressions
176 deListExpr cons nil (HsDoOut ListComp stmts _ _ _ result_ty src_loc)
177 = deListComp stmts cons nil
179 deListExpr cons nil (HsVar map, _, [f,xs])
180 | goodInst var mapIdKey = dsExpr f `thenDs` \ f' ->
181 feComposeWith cons f' `thenDs` \ cons' ->
183 deListExpr xs cons' nil
186 data HsExprForm = GoodForm What [Type] [TypecheckedHsExpr]
189 data What = HsMap | HsConcat | HsFilter | HsZip | HsFoldr
191 analyseListProducer (HsVar v) ty_args val_args
192 | good_inst mapIdKey 2 = GoodForm HsMap ty_args val_args
193 | good_inst concatIdKey 1 = GoodForm HsConcat ty_args val_args
194 | good_inst filterIdKey 2 = GoodForm HsFilter ty_args val_args
195 | good_id zipIdKey 2 = GoodForm HsZip ty_args val_args
198 good_inst key arity = isInstIdOf key v && result_is_list && n_args == arity
199 good_id key arity = getUnique v == key && result_is_list && n_args == arity
202 n_args = length val_args
204 result_is_list = resultTyIsList (idType v) ty_args val_args
206 resultTyIsList ty ty_args val_args
209 go1 ty (_:tys) = case splitForAllTy_maybe ty of
210 Just (_,ty) -> go1 ty tys
212 go1 ty [] = go2 ty val_args
214 go2 ty (_:args) = case splitFunTy_maybe of
215 Just (_,ty) -> go2 ty args
218 go2 ty [] = case splitTyConApp_maybe of
219 Just (tycon, [_]) | tycon == listTyCon -> True
224 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
226 deListComp [ReturnStmt expr] list -- Figure 7.4, SLPJ, p 135, rule C above
227 = dsExpr expr `thenDs` \ core_expr ->
228 returnDs (mkConApp consDataCon [Type (coreExprType core_expr), core_expr, list])
230 deListComp (GuardStmt guard locn : quals) list -- rule B above
231 = dsExpr guard `thenDs` \ core_guard ->
232 deListComp quals list `thenDs` \ core_rest ->
233 returnDs (mkIfThenElse core_guard core_rest list)
235 -- [e | let B, qs] = let B in [e | qs]
236 deListComp (LetStmt binds : quals) list
237 = deListComp quals list `thenDs` \ core_rest ->
238 dsLet binds core_rest
240 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
241 = dsExpr list1 `thenDs` \ core_list1 ->
243 u3_ty@u1_ty = coreExprType core_list1 -- two names, same thing
245 -- u1_ty is a [alpha] type, and u2_ty = alpha
246 u2_ty = outPatType pat
248 res_ty = coreExprType core_list2
249 h_ty = u1_ty `mkFunTy` res_ty
251 newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] `thenDs` \ [h, u1, u2, u3] ->
253 -- the "fail" value ...
255 core_fail = App (Var h) (Var u3)
256 letrec_body = App (Var h) core_list1
258 deListComp quals core_fail `thenDs` \ rest_expr ->
259 matchSimply (Var u2) ListCompMatch pat
260 rest_expr core_fail `thenDs` \ core_match ->
263 Case (Var u1) u1 [(DataCon nilDataCon, [], core_list2),
264 (DataCon consDataCon, [u2, u3], core_match)]
266 returnDs (Let (Rec [(h, rhs)]) letrec_body)
269 %************************************************************************
271 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
273 %************************************************************************
275 @dfListComp@ are the rules used with foldr/build turned on:
277 TE < [ e | ] >> c n = c e n
278 TE << [ e | b , q ] >> c n = if b then TE << [ e | q ] >> c n else n
279 TE << [ e | p <- l , q ] c n = foldr
280 (\ TE << p >> b -> TE << [ e | q ] >> c b
284 dfListComp :: Type -> Id -- 'c'; its type and id
285 -> Type -> Id -- 'n'; its type and id
286 -> [TypecheckedStmt] -- the rest of the qual's
289 dfListComp c_ty c_id n_ty n_id [ReturnStmt expr]
290 = dsExpr expr `thenDs` \ core_expr ->
291 returnDs (mkApps (Var c_id) [core_expr, Var n_id])
293 dfListComp c_ty c_id n_ty n_id (GuardStmt guard locn : quals)
294 = dsExpr guard `thenDs` \ core_guard ->
295 dfListComp c_ty c_id n_ty n_id quals `thenDs` \ core_rest ->
296 returnDs (mkIfThenElse core_guard core_rest (Var n_id))
298 dfListComp c_ty c_id n_ty n_id (LetStmt binds : quals)
299 -- new in 1.3, local bindings
300 = dfListComp c_ty c_id n_ty n_id quals `thenDs` \ core_rest ->
301 dsLet binds core_rest
303 dfListComp c_ty c_id n_ty n_id (BindStmt pat list1 locn : quals)
304 -- evaluate the two lists
305 = dsExpr list1 `thenDs` \ core_list1 ->
307 -- find the required type
309 let p_ty = outPatType pat
310 b_ty = n_ty -- alias b_ty to n_ty
311 fn_ty = mkFunTys [p_ty, b_ty] b_ty
312 lst_ty = coreExprType core_list1
315 -- create some new local id's
317 newSysLocalsDs [b_ty,p_ty,fn_ty,lst_ty] `thenDs` \ [b,p,fn,lst] ->
319 -- build rest of the comprehesion
321 dfListComp c_ty c_id b_ty b quals `thenDs` \ core_rest ->
322 -- build the pattern match
324 matchSimply (Var p) ListCompMatch pat core_rest (Var b) `thenDs` \ core_expr ->
326 -- now build the outermost foldr, and return
330 [NonRec fn (mkLams [p, b] core_expr),
331 NonRec lst core_list1]
332 (mkFoldr p_ty n_ty fn n_id lst)
337 @mkBuild@ is sugar for building a build!
339 @mkbuild ty tv c n e@ $Rightarrow$ @build ty (/\ tv -> \ c n -> e)@
340 @ty@ is the type of the list.
341 @tv@ is always a new type variable.
342 @c,n@ are Id's for the abstract cons and nil, @g@ for let binding the argument argument.
345 v :: (\/ b . (a -> b -> b) -> b -> b) -> [a]
346 -- \/ a . (\/ b . (a -> b -> b) -> b -> b) -> [a]
347 @e@ is the object right inside the @build@
355 -> CoreExpr -- template
356 -> CoreExpr -- template
358 mkBuild ty tv c n g expr
359 = Let (NonRec g (mkLams [tv, c,n] expr))
360 (mkApps (Var buildId) [Type ty, Var g])
362 buildId = error "DsListComp: buildId"
365 = mkApps (mkTyApps (Var foldrId) [a,b]) [Var f, Var z, Var xs]