TTG3 Combined Step 1 and 3 for Trees That Grow
[ghc.git] / compiler / deSugar / DsListComp.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4
5
6 Desugaring list comprehensions, monad comprehensions and array comprehensions
7 -}
8
9 {-# LANGUAGE CPP, NamedFieldPuns #-}
10 {-# LANGUAGE TypeFamilies #-}
11
12 module DsListComp ( dsListComp, dsPArrComp, dsMonadComp ) where
13
14 #include "HsVersions.h"
15
16 import GhcPrelude
17
18 import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLExprNoLP, dsLocalBinds, dsSyntaxExpr )
19
20 import HsSyn
21 import TcHsSyn
22 import CoreSyn
23 import MkCore
24
25 import DsMonad -- the monadery used in the desugarer
26 import DsUtils
27
28 import DynFlags
29 import CoreUtils
30 import Id
31 import Type
32 import TysWiredIn
33 import Match
34 import PrelNames
35 import SrcLoc
36 import Outputable
37 import TcType
38 import ListSetOps( getNth )
39 import Util
40
41 {-
42 List comprehensions may be desugared in one of two ways: ``ordinary''
43 (as you would expect if you read SLPJ's book) and ``with foldr/build
44 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
45
46 There will be at least one ``qualifier'' in the input.
47 -}
48
49 dsListComp :: [ExprLStmt GhcTc]
50 -> Type -- Type of entire list
51 -> DsM CoreExpr
52 dsListComp lquals res_ty = do
53 dflags <- getDynFlags
54 let quals = map unLoc lquals
55 elt_ty = case tcTyConAppArgs res_ty of
56 [elt_ty] -> elt_ty
57 _ -> pprPanic "dsListComp" (ppr res_ty $$ ppr lquals)
58
59 if not (gopt Opt_EnableRewriteRules dflags) || gopt Opt_IgnoreInterfacePragmas dflags
60 -- Either rules are switched off, or we are ignoring what there are;
61 -- Either way foldr/build won't happen, so use the more efficient
62 -- Wadler-style desugaring
63 || isParallelComp quals
64 -- Foldr-style desugaring can't handle parallel list comprehensions
65 then deListComp quals (mkNilExpr elt_ty)
66 else mkBuildExpr elt_ty (\(c, _) (n, _) -> dfListComp c n quals)
67 -- Foldr/build should be enabled, so desugar
68 -- into foldrs and builds
69
70 where
71 -- We must test for ParStmt anywhere, not just at the head, because an extension
72 -- to list comprehensions would be to add brackets to specify the associativity
73 -- of qualifier lists. This is really easy to do by adding extra ParStmts into the
74 -- mix of possibly a single element in length, so we do this to leave the possibility open
75 isParallelComp = any isParallelStmt
76
77 isParallelStmt (ParStmt {}) = True
78 isParallelStmt _ = False
79
80
81 -- This function lets you desugar a inner list comprehension and a list of the binders
82 -- of that comprehension that we need in the outer comprehension into such an expression
83 -- and the type of the elements that it outputs (tuples of binders)
84 dsInnerListComp :: (ParStmtBlock GhcTc GhcTc) -> DsM (CoreExpr, Type)
85 dsInnerListComp (ParStmtBlock _ stmts bndrs _)
86 = do { let bndrs_tuple_type = mkBigCoreVarTupTy bndrs
87 list_ty = mkListTy bndrs_tuple_type
88
89 -- really use original bndrs below!
90 ; expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTupId bndrs)]) list_ty
91
92 ; return (expr, bndrs_tuple_type) }
93 dsInnerListComp (XParStmtBlock{}) = panic "dsInnerListComp"
94
95 -- This function factors out commonality between the desugaring strategies for GroupStmt.
96 -- Given such a statement it gives you back an expression representing how to compute the transformed
97 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
98 dsTransStmt :: ExprStmt GhcTc -> DsM (CoreExpr, LPat GhcTc)
99 dsTransStmt (TransStmt { trS_form = form, trS_stmts = stmts, trS_bndrs = binderMap
100 , trS_by = by, trS_using = using }) = do
101 let (from_bndrs, to_bndrs) = unzip binderMap
102
103 let from_bndrs_tys = map idType from_bndrs
104 to_bndrs_tys = map idType to_bndrs
105
106 to_bndrs_tup_ty = mkBigCoreTupTy to_bndrs_tys
107
108 -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
109 (expr', from_tup_ty) <- dsInnerListComp (ParStmtBlock noExt stmts
110 from_bndrs noSyntaxExpr)
111
112 -- Work out what arguments should be supplied to that expression: i.e. is an extraction
113 -- function required? If so, create that desugared function and add to arguments
114 usingExpr' <- dsLExpr using
115 usingArgs' <- case by of
116 Nothing -> return [expr']
117 Just by_e -> do { by_e' <- dsLExpr by_e
118 ; lam' <- matchTuple from_bndrs by_e'
119 ; return [lam', expr'] }
120
121 -- Create an unzip function for the appropriate arity and element types and find "map"
122 unzip_stuff' <- mkUnzipBind form from_bndrs_tys
123 map_id <- dsLookupGlobalId mapName
124
125 -- Generate the expressions to build the grouped list
126 let -- First we apply the grouping function to the inner list
127 inner_list_expr' = mkApps usingExpr' usingArgs'
128 -- Then we map our "unzip" across it to turn the lists of tuples into tuples of lists
129 -- We make sure we instantiate the type variable "a" to be a list of "from" tuples and
130 -- the "b" to be a tuple of "to" lists!
131 -- Then finally we bind the unzip function around that expression
132 bound_unzipped_inner_list_expr'
133 = case unzip_stuff' of
134 Nothing -> inner_list_expr'
135 Just (unzip_fn', unzip_rhs') ->
136 Let (Rec [(unzip_fn', unzip_rhs')]) $
137 mkApps (Var map_id) $
138 [ Type (mkListTy from_tup_ty)
139 , Type to_bndrs_tup_ty
140 , Var unzip_fn'
141 , inner_list_expr' ]
142
143 dsNoLevPoly (tcFunResultTyN (length usingArgs') (exprType usingExpr'))
144 (text "In the result of a" <+> quotes (text "using") <+> text "function:" <+> ppr using)
145
146 -- Build a pattern that ensures the consumer binds into the NEW binders,
147 -- which hold lists rather than single values
148 let pat = mkBigLHsVarPatTupId to_bndrs -- NB: no '!
149 return (bound_unzipped_inner_list_expr', pat)
150
151 dsTransStmt _ = panic "dsTransStmt: Not given a TransStmt"
152
153 {-
154 ************************************************************************
155 * *
156 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
157 * *
158 ************************************************************************
159
160 Just as in Phil's chapter~7 in SLPJ, using the rules for
161 optimally-compiled list comprehensions. This is what Kevin followed
162 as well, and I quite happily do the same. The TQ translation scheme
163 transforms a list of qualifiers (either boolean expressions or
164 generators) into a single expression which implements the list
165 comprehension. Because we are generating 2nd-order polymorphic
166 lambda-calculus, calls to NIL and CONS must be applied to a type
167 argument, as well as their usual value arguments.
168 \begin{verbatim}
169 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
170
171 (Rule C)
172 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
173
174 (Rule B)
175 TQ << [ e | b , qs ] ++ L >> =
176 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
177
178 (Rule A')
179 TQ << [ e | p <- L1, qs ] ++ L2 >> =
180 letrec
181 h = \ u1 ->
182 case u1 of
183 [] -> TE << L2 >>
184 (u2 : u3) ->
185 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
186 [] (h u3)
187 in
188 h ( TE << L1 >> )
189
190 "h", "u1", "u2", and "u3" are new variables.
191 \end{verbatim}
192
193 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
194 is the TE translation scheme. Note that we carry around the @L@ list
195 already desugared. @dsListComp@ does the top TE rule mentioned above.
196
197 To the above, we add an additional rule to deal with parallel list
198 comprehensions. The translation goes roughly as follows:
199 [ e | p1 <- e11, let v1 = e12, p2 <- e13
200 | q1 <- e21, let v2 = e22, q2 <- e23]
201 =>
202 [ e | ((x1, .., xn), (y1, ..., ym)) <-
203 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
204 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
205 where (x1, .., xn) are the variables bound in p1, v1, p2
206 (y1, .., ym) are the variables bound in q1, v2, q2
207
208 In the translation below, the ParStmt branch translates each parallel branch
209 into a sub-comprehension, and desugars each independently. The resulting lists
210 are fed to a zip function, we create a binding for all the variables bound in all
211 the comprehensions, and then we hand things off the the desugarer for bindings.
212 The zip function is generated here a) because it's small, and b) because then we
213 don't have to deal with arbitrary limits on the number of zip functions in the
214 prelude, nor which library the zip function came from.
215 The introduced tuples are Boxed, but only because I couldn't get it to work
216 with the Unboxed variety.
217 -}
218
219 deListComp :: [ExprStmt GhcTc] -> CoreExpr -> DsM CoreExpr
220
221 deListComp [] _ = panic "deListComp"
222
223 deListComp (LastStmt body _ _ : quals) list
224 = -- Figure 7.4, SLPJ, p 135, rule C above
225 ASSERT( null quals )
226 do { core_body <- dsLExpr body
227 ; return (mkConsExpr (exprType core_body) core_body list) }
228
229 -- Non-last: must be a guard
230 deListComp (BodyStmt guard _ _ _ : quals) list = do -- rule B above
231 core_guard <- dsLExpr guard
232 core_rest <- deListComp quals list
233 return (mkIfThenElse core_guard core_rest list)
234
235 -- [e | let B, qs] = let B in [e | qs]
236 deListComp (LetStmt binds : quals) list = do
237 core_rest <- deListComp quals list
238 dsLocalBinds binds core_rest
239
240 deListComp (stmt@(TransStmt {}) : quals) list = do
241 (inner_list_expr, pat) <- dsTransStmt stmt
242 deBindComp pat inner_list_expr quals list
243
244 deListComp (BindStmt pat list1 _ _ _ : quals) core_list2 = do -- rule A' above
245 core_list1 <- dsLExprNoLP list1
246 deBindComp pat core_list1 quals core_list2
247
248 deListComp (ParStmt stmtss_w_bndrs _ _ _ : quals) list
249 = do { exps_and_qual_tys <- mapM dsInnerListComp stmtss_w_bndrs
250 ; let (exps, qual_tys) = unzip exps_and_qual_tys
251
252 ; (zip_fn, zip_rhs) <- mkZipBind qual_tys
253
254 -- Deal with [e | pat <- zip l1 .. ln] in example above
255 ; deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
256 quals list }
257 where
258 bndrs_s = [bs | ParStmtBlock _ _ bs _ <- stmtss_w_bndrs]
259
260 -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
261 pat = mkBigLHsPatTupId pats
262 pats = map mkBigLHsVarPatTupId bndrs_s
263
264 deListComp (RecStmt {} : _) _ = panic "deListComp RecStmt"
265
266 deListComp (ApplicativeStmt {} : _) _ =
267 panic "deListComp ApplicativeStmt"
268
269 deBindComp :: OutPat GhcTc
270 -> CoreExpr
271 -> [ExprStmt GhcTc]
272 -> CoreExpr
273 -> DsM (Expr Id)
274 deBindComp pat core_list1 quals core_list2 = do
275 let u3_ty@u1_ty = exprType core_list1 -- two names, same thing
276
277 -- u1_ty is a [alpha] type, and u2_ty = alpha
278 let u2_ty = hsLPatType pat
279
280 let res_ty = exprType core_list2
281 h_ty = u1_ty `mkFunTy` res_ty
282
283 -- no levity polymorphism here, as list comprehensions don't work
284 -- with RebindableSyntax. NB: These are *not* monad comps.
285 [h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]
286
287 -- the "fail" value ...
288 let
289 core_fail = App (Var h) (Var u3)
290 letrec_body = App (Var h) core_list1
291
292 rest_expr <- deListComp quals core_fail
293 core_match <- matchSimply (Var u2) (StmtCtxt ListComp) pat rest_expr core_fail
294
295 let
296 rhs = Lam u1 $
297 Case (Var u1) u1 res_ty
298 [(DataAlt nilDataCon, [], core_list2),
299 (DataAlt consDataCon, [u2, u3], core_match)]
300 -- Increasing order of tag
301
302 return (Let (Rec [(h, rhs)]) letrec_body)
303
304 {-
305 ************************************************************************
306 * *
307 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
308 * *
309 ************************************************************************
310
311 @dfListComp@ are the rules used with foldr/build turned on:
312
313 \begin{verbatim}
314 TE[ e | ] c n = c e n
315 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
316 TE[ e | p <- l , q ] c n = let
317 f = \ x b -> case x of
318 p -> TE[ e | q ] c b
319 _ -> b
320 in
321 foldr f n l
322 \end{verbatim}
323 -}
324
325 dfListComp :: Id -> Id -- 'c' and 'n'
326 -> [ExprStmt GhcTc] -- the rest of the qual's
327 -> DsM CoreExpr
328
329 dfListComp _ _ [] = panic "dfListComp"
330
331 dfListComp c_id n_id (LastStmt body _ _ : quals)
332 = ASSERT( null quals )
333 do { core_body <- dsLExprNoLP body
334 ; return (mkApps (Var c_id) [core_body, Var n_id]) }
335
336 -- Non-last: must be a guard
337 dfListComp c_id n_id (BodyStmt guard _ _ _ : quals) = do
338 core_guard <- dsLExpr guard
339 core_rest <- dfListComp c_id n_id quals
340 return (mkIfThenElse core_guard core_rest (Var n_id))
341
342 dfListComp c_id n_id (LetStmt binds : quals) = do
343 -- new in 1.3, local bindings
344 core_rest <- dfListComp c_id n_id quals
345 dsLocalBinds binds core_rest
346
347 dfListComp c_id n_id (stmt@(TransStmt {}) : quals) = do
348 (inner_list_expr, pat) <- dsTransStmt stmt
349 -- Anyway, we bind the newly grouped list via the generic binding function
350 dfBindComp c_id n_id (pat, inner_list_expr) quals
351
352 dfListComp c_id n_id (BindStmt pat list1 _ _ _ : quals) = do
353 -- evaluate the two lists
354 core_list1 <- dsLExpr list1
355
356 -- Do the rest of the work in the generic binding builder
357 dfBindComp c_id n_id (pat, core_list1) quals
358
359 dfListComp _ _ (ParStmt {} : _) = panic "dfListComp ParStmt"
360 dfListComp _ _ (RecStmt {} : _) = panic "dfListComp RecStmt"
361 dfListComp _ _ (ApplicativeStmt {} : _) =
362 panic "dfListComp ApplicativeStmt"
363
364 dfBindComp :: Id -> Id -- 'c' and 'n'
365 -> (LPat GhcTc, CoreExpr)
366 -> [ExprStmt GhcTc] -- the rest of the qual's
367 -> DsM CoreExpr
368 dfBindComp c_id n_id (pat, core_list1) quals = do
369 -- find the required type
370 let x_ty = hsLPatType pat
371 let b_ty = idType n_id
372
373 -- create some new local id's
374 b <- newSysLocalDs b_ty
375 x <- newSysLocalDs x_ty
376
377 -- build rest of the comprehesion
378 core_rest <- dfListComp c_id b quals
379
380 -- build the pattern match
381 core_expr <- matchSimply (Var x) (StmtCtxt ListComp)
382 pat core_rest (Var b)
383
384 -- now build the outermost foldr, and return
385 mkFoldrExpr x_ty b_ty (mkLams [x, b] core_expr) (Var n_id) core_list1
386
387 {-
388 ************************************************************************
389 * *
390 \subsection[DsFunGeneration]{Generation of zip/unzip functions for use in desugaring}
391 * *
392 ************************************************************************
393 -}
394
395 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
396 -- mkZipBind [t1, t2]
397 -- = (zip, \as1:[t1] as2:[t2]
398 -- -> case as1 of
399 -- [] -> []
400 -- (a1:as'1) -> case as2 of
401 -- [] -> []
402 -- (a2:as'2) -> (a1, a2) : zip as'1 as'2)]
403
404 mkZipBind elt_tys = do
405 ass <- mapM newSysLocalDs elt_list_tys
406 as' <- mapM newSysLocalDs elt_tys
407 as's <- mapM newSysLocalDs elt_list_tys
408
409 zip_fn <- newSysLocalDs zip_fn_ty
410
411 let inner_rhs = mkConsExpr elt_tuple_ty
412 (mkBigCoreVarTup as')
413 (mkVarApps (Var zip_fn) as's)
414 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
415
416 return (zip_fn, mkLams ass zip_body)
417 where
418 elt_list_tys = map mkListTy elt_tys
419 elt_tuple_ty = mkBigCoreTupTy elt_tys
420 elt_tuple_list_ty = mkListTy elt_tuple_ty
421
422 zip_fn_ty = mkFunTys elt_list_tys elt_tuple_list_ty
423
424 mk_case (as, a', as') rest
425 = Case (Var as) as elt_tuple_list_ty
426 [(DataAlt nilDataCon, [], mkNilExpr elt_tuple_ty),
427 (DataAlt consDataCon, [a', as'], rest)]
428 -- Increasing order of tag
429
430
431 mkUnzipBind :: TransForm -> [Type] -> DsM (Maybe (Id, CoreExpr))
432 -- mkUnzipBind [t1, t2]
433 -- = (unzip, \ys :: [(t1, t2)] -> foldr (\ax :: (t1, t2) axs :: ([t1], [t2])
434 -- -> case ax of
435 -- (x1, x2) -> case axs of
436 -- (xs1, xs2) -> (x1 : xs1, x2 : xs2))
437 -- ([], [])
438 -- ys)
439 --
440 -- We use foldr here in all cases, even if rules are turned off, because we may as well!
441 mkUnzipBind ThenForm _
442 = return Nothing -- No unzipping for ThenForm
443 mkUnzipBind _ elt_tys
444 = do { ax <- newSysLocalDs elt_tuple_ty
445 ; axs <- newSysLocalDs elt_list_tuple_ty
446 ; ys <- newSysLocalDs elt_tuple_list_ty
447 ; xs <- mapM newSysLocalDs elt_tys
448 ; xss <- mapM newSysLocalDs elt_list_tys
449
450 ; unzip_fn <- newSysLocalDs unzip_fn_ty
451
452 ; [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
453
454 ; let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
455 concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
456 tupled_concat_expression = mkBigCoreTup concat_expressions
457
458 folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
459 folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
460 folder_body = mkLams [ax, axs] folder_body_outer_case
461
462 ; unzip_body <- mkFoldrExpr elt_tuple_ty elt_list_tuple_ty folder_body nil_tuple (Var ys)
463 ; return (Just (unzip_fn, mkLams [ys] unzip_body)) }
464 where
465 elt_tuple_ty = mkBigCoreTupTy elt_tys
466 elt_tuple_list_ty = mkListTy elt_tuple_ty
467 elt_list_tys = map mkListTy elt_tys
468 elt_list_tuple_ty = mkBigCoreTupTy elt_list_tys
469
470 unzip_fn_ty = elt_tuple_list_ty `mkFunTy` elt_list_tuple_ty
471
472 mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
473
474 {-
475 ************************************************************************
476 * *
477 \subsection[DsPArrComp]{Desugaring of array comprehensions}
478 * *
479 ************************************************************************
480 -}
481
482 -- entry point for desugaring a parallel array comprehension
483 --
484 -- [:e | qss:] = <<[:e | qss:]>> () [:():]
485 --
486 dsPArrComp :: [ExprStmt GhcTc]
487 -> DsM CoreExpr
488
489 -- Special case for parallel comprehension
490 dsPArrComp (ParStmt qss _ _ _ : quals) = dePArrParComp qss quals
491
492 -- Special case for simple generators:
493 --
494 -- <<[:e' | p <- e, qs:]>> = <<[: e' | qs :]>> p e
495 --
496 -- if matching again p cannot fail, or else
497 --
498 -- <<[:e' | p <- e, qs:]>> =
499 -- <<[:e' | qs:]>> p (filterP (\x -> case x of {p -> True; _ -> False}) e)
500 --
501 dsPArrComp (BindStmt p e _ _ _ : qs) = do
502 filterP <- dsDPHBuiltin filterPVar
503 ce <- dsLExprNoLP e
504 let ety'ce = parrElemType ce
505 false = Var falseDataConId
506 true = Var trueDataConId
507 v <- newSysLocalDs ety'ce
508 pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
509 let gen | isIrrefutableHsPat p = ce
510 | otherwise = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
511 dePArrComp qs p gen
512
513 dsPArrComp qs = do -- no ParStmt in `qs'
514 sglP <- dsDPHBuiltin singletonPVar
515 let unitArray = mkApps (Var sglP) [Type unitTy, mkCoreTup []]
516 dePArrComp qs (noLoc $ WildPat unitTy) unitArray
517
518
519
520 -- the work horse
521 --
522 dePArrComp :: [ExprStmt GhcTc]
523 -> LPat GhcTc -- the current generator pattern
524 -> CoreExpr -- the current generator expression
525 -> DsM CoreExpr
526
527 dePArrComp [] _ _ = panic "dePArrComp"
528
529 --
530 -- <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
531 --
532 dePArrComp (LastStmt e' _ _ : quals) pa cea
533 = ASSERT( null quals )
534 do { mapP <- dsDPHBuiltin mapPVar
535 ; let ty = parrElemType cea
536 ; (clam, ty'e') <- deLambda ty pa e'
537 ; return $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea] }
538 --
539 -- <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
540 --
541 dePArrComp (BodyStmt b _ _ _ : qs) pa cea = do
542 filterP <- dsDPHBuiltin filterPVar
543 let ty = parrElemType cea
544 (clam,_) <- deLambda ty pa b
545 dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
546
547 --
548 -- <<[:e' | p <- e, qs:]>> pa ea =
549 -- let ef = \pa -> e
550 -- in
551 -- <<[:e' | qs:]>> (pa, p) (crossMap ea ef)
552 --
553 -- if matching again p cannot fail, or else
554 --
555 -- <<[:e' | p <- e, qs:]>> pa ea =
556 -- let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
557 -- in
558 -- <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
559 --
560 dePArrComp (BindStmt p e _ _ _ : qs) pa cea = do
561 filterP <- dsDPHBuiltin filterPVar
562 crossMapP <- dsDPHBuiltin crossMapPVar
563 ce <- dsLExpr e
564 let ety'cea = parrElemType cea
565 ety'ce = parrElemType ce
566 false = Var falseDataConId
567 true = Var trueDataConId
568 v <- newSysLocalDs ety'ce
569 pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
570 let cef | isIrrefutableHsPat p = ce
571 | otherwise = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
572 (clam, _) <- mkLambda ety'cea pa cef
573 let ety'cef = ety'ce -- filter doesn't change the element type
574 pa' = mkLHsPatTup [pa, p]
575
576 dePArrComp qs pa' (mkApps (Var crossMapP)
577 [Type ety'cea, Type ety'cef, cea, clam])
578 --
579 -- <<[:e' | let ds, qs:]>> pa ea =
580 -- <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
581 -- (mapP (\v@pa -> let ds in (v, (x_1, ..., x_n))) ea)
582 -- where
583 -- {x_1, ..., x_n} = DV (ds) -- Defined Variables
584 --
585 dePArrComp (LetStmt lds@(L _ ds) : qs) pa cea = do
586 mapP <- dsDPHBuiltin mapPVar
587 let xs = collectLocalBinders ds
588 ty'cea = parrElemType cea
589 v <- newSysLocalDs ty'cea
590 clet <- dsLocalBinds lds (mkCoreTup (map Var xs))
591 let'v <- newSysLocalDs (exprType clet)
592 let projBody = mkCoreLet (NonRec let'v clet) $
593 mkCoreTup [Var v, Var let'v]
594 errTy = exprType projBody
595 errMsg = text "DsListComp.dePArrComp: internal error!"
596 cerr <- mkErrorAppDs pAT_ERROR_ID errTy errMsg
597 ccase <- matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr
598 let pa' = mkLHsPatTup [pa, mkLHsPatTup (map nlVarPat xs)]
599 proj = mkLams [v] ccase
600 dePArrComp qs pa' (mkApps (Var mapP)
601 [Type ty'cea, Type errTy, proj, cea])
602 --
603 -- The parser guarantees that parallel comprehensions can only appear as
604 -- singleton qualifier lists, which we already special case in the caller.
605 -- So, encountering one here is a bug.
606 --
607 dePArrComp (ParStmt {} : _) _ _ =
608 panic "DsListComp.dePArrComp: malformed comprehension AST: ParStmt"
609 dePArrComp (TransStmt {} : _) _ _ = panic "DsListComp.dePArrComp: TransStmt"
610 dePArrComp (RecStmt {} : _) _ _ = panic "DsListComp.dePArrComp: RecStmt"
611 dePArrComp (ApplicativeStmt {} : _) _ _ =
612 panic "DsListComp.dePArrComp: ApplicativeStmt"
613
614 -- <<[:e' | qs | qss:]>> pa ea =
615 -- <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
616 -- (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
617 -- where
618 -- {x_1, ..., x_n} = DV (qs)
619 --
620 dePArrParComp :: [ParStmtBlock GhcTc GhcTc] -> [ExprStmt GhcTc] -> DsM CoreExpr
621 dePArrParComp qss quals = do
622 (pQss, ceQss) <- deParStmt qss
623 dePArrComp quals pQss ceQss
624 where
625 deParStmt [] =
626 -- empty parallel statement lists have no source representation
627 panic "DsListComp.dePArrComp: Empty parallel list comprehension"
628 deParStmt (ParStmtBlock _ qs xs _:qss) = do -- first statement
629 let res_expr = mkLHsVarTuple xs
630 cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
631 parStmts qss (mkLHsVarPatTup xs) cqs
632 deParStmt (XParStmtBlock{}:_) = panic "dePArrParComp"
633 ---
634 parStmts [] pa cea = return (pa, cea)
635 parStmts (ParStmtBlock _ qs xs _:qss) pa cea = do
636 -- subsequent statements (zip'ed)
637 zipP <- dsDPHBuiltin zipPVar
638 let pa' = mkLHsPatTup [pa, mkLHsVarPatTup xs]
639 ty'cea = parrElemType cea
640 res_expr = mkLHsVarTuple xs
641 cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
642 let ty'cqs = parrElemType cqs
643 cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
644 parStmts qss pa' cea'
645 parStmts (XParStmtBlock{}:_) _ _ = panic "dePArrParComp"
646
647 -- generate Core corresponding to `\p -> e'
648 --
649 deLambda :: Type -- type of the argument (not levity-polymorphic)
650 -> LPat GhcTc -- argument pattern
651 -> LHsExpr GhcTc -- body
652 -> DsM (CoreExpr, Type)
653 deLambda ty p e =
654 mkLambda ty p =<< dsLExpr e
655
656 -- generate Core for a lambda pattern match, where the body is already in Core
657 --
658 mkLambda :: Type -- type of the argument (not levity-polymorphic)
659 -> LPat GhcTc -- argument pattern
660 -> CoreExpr -- desugared body
661 -> DsM (CoreExpr, Type)
662 mkLambda ty p ce = do
663 v <- newSysLocalDs ty
664 let errMsg = text "DsListComp.deLambda: internal error!"
665 ce'ty = exprType ce
666 cerr <- mkErrorAppDs pAT_ERROR_ID ce'ty errMsg
667 res <- matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr
668 return (mkLams [v] res, ce'ty)
669
670 -- obtain the element type of the parallel array produced by the given Core
671 -- expression
672 --
673 parrElemType :: CoreExpr -> Type
674 parrElemType e =
675 case splitTyConApp_maybe (exprType e) of
676 Just (tycon, [ty]) | tycon == parrTyCon -> ty
677 _ -> panic
678 "DsListComp.parrElemType: not a parallel array type"
679
680 -- Translation for monad comprehensions
681
682 -- Entry point for monad comprehension desugaring
683 dsMonadComp :: [ExprLStmt GhcTc] -> DsM CoreExpr
684 dsMonadComp stmts = dsMcStmts stmts
685
686 dsMcStmts :: [ExprLStmt GhcTc] -> DsM CoreExpr
687 dsMcStmts [] = panic "dsMcStmts"
688 dsMcStmts (L loc stmt : lstmts) = putSrcSpanDs loc (dsMcStmt stmt lstmts)
689
690 ---------------
691 dsMcStmt :: ExprStmt GhcTc -> [ExprLStmt GhcTc] -> DsM CoreExpr
692
693 dsMcStmt (LastStmt body _ ret_op) stmts
694 = ASSERT( null stmts )
695 do { body' <- dsLExpr body
696 ; dsSyntaxExpr ret_op [body'] }
697
698 -- [ .. | let binds, stmts ]
699 dsMcStmt (LetStmt binds) stmts
700 = do { rest <- dsMcStmts stmts
701 ; dsLocalBinds binds rest }
702
703 -- [ .. | a <- m, stmts ]
704 dsMcStmt (BindStmt pat rhs bind_op fail_op bind_ty) stmts
705 = do { rhs' <- dsLExpr rhs
706 ; dsMcBindStmt pat rhs' bind_op fail_op bind_ty stmts }
707
708 -- Apply `guard` to the `exp` expression
709 --
710 -- [ .. | exp, stmts ]
711 --
712 dsMcStmt (BodyStmt exp then_exp guard_exp _) stmts
713 = do { exp' <- dsLExpr exp
714 ; rest <- dsMcStmts stmts
715 ; guard_exp' <- dsSyntaxExpr guard_exp [exp']
716 ; dsSyntaxExpr then_exp [guard_exp', rest] }
717
718 -- Group statements desugar like this:
719 --
720 -- [| (q, then group by e using f); rest |]
721 -- ---> f {qt} (\qv -> e) [| q; return qv |] >>= \ n_tup ->
722 -- case unzip n_tup of qv' -> [| rest |]
723 --
724 -- where variables (v1:t1, ..., vk:tk) are bound by q
725 -- qv = (v1, ..., vk)
726 -- qt = (t1, ..., tk)
727 -- (>>=) :: m2 a -> (a -> m3 b) -> m3 b
728 -- f :: forall a. (a -> t) -> m1 a -> m2 (n a)
729 -- n_tup :: n qt
730 -- unzip :: n qt -> (n t1, ..., n tk) (needs Functor n)
731
732 dsMcStmt (TransStmt { trS_stmts = stmts, trS_bndrs = bndrs
733 , trS_by = by, trS_using = using
734 , trS_ret = return_op, trS_bind = bind_op
735 , trS_bind_arg_ty = n_tup_ty' -- n (a,b,c)
736 , trS_fmap = fmap_op, trS_form = form }) stmts_rest
737 = do { let (from_bndrs, to_bndrs) = unzip bndrs
738
739 ; let from_bndr_tys = map idType from_bndrs -- Types ty
740
741
742 -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
743 ; expr' <- dsInnerMonadComp stmts from_bndrs return_op
744
745 -- Work out what arguments should be supplied to that expression: i.e. is an extraction
746 -- function required? If so, create that desugared function and add to arguments
747 ; usingExpr' <- dsLExpr using
748 ; usingArgs' <- case by of
749 Nothing -> return [expr']
750 Just by_e -> do { by_e' <- dsLExpr by_e
751 ; lam' <- matchTuple from_bndrs by_e'
752 ; return [lam', expr'] }
753
754 -- Generate the expressions to build the grouped list
755 -- Build a pattern that ensures the consumer binds into the NEW binders,
756 -- which hold monads rather than single values
757 ; let tup_n_ty' = mkBigCoreVarTupTy to_bndrs
758
759 ; body <- dsMcStmts stmts_rest
760 ; n_tup_var' <- newSysLocalDsNoLP n_tup_ty'
761 ; tup_n_var' <- newSysLocalDs tup_n_ty'
762 ; tup_n_expr' <- mkMcUnzipM form fmap_op n_tup_var' from_bndr_tys
763 ; us <- newUniqueSupply
764 ; let rhs' = mkApps usingExpr' usingArgs'
765 body' = mkTupleCase us to_bndrs body tup_n_var' tup_n_expr'
766
767 ; dsSyntaxExpr bind_op [rhs', Lam n_tup_var' body'] }
768
769 -- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel
770 -- statements, for example:
771 --
772 -- [ body | qs1 | qs2 | qs3 ]
773 -- -> [ body | (bndrs1, (bndrs2, bndrs3))
774 -- <- [bndrs1 | qs1] `mzip` ([bndrs2 | qs2] `mzip` [bndrs3 | qs3]) ]
775 --
776 -- where `mzip` has type
777 -- mzip :: forall a b. m a -> m b -> m (a,b)
778 -- NB: we need a polymorphic mzip because we call it several times
779
780 dsMcStmt (ParStmt blocks mzip_op bind_op bind_ty) stmts_rest
781 = do { exps_w_tys <- mapM ds_inner blocks -- Pairs (exp :: m ty, ty)
782 ; mzip_op' <- dsExpr mzip_op
783
784 ; let -- The pattern variables
785 pats = [ mkBigLHsVarPatTupId bs | ParStmtBlock _ _ bs _ <- blocks]
786 -- Pattern with tuples of variables
787 -- [v1,v2,v3] => (v1, (v2, v3))
788 pat = foldr1 (\p1 p2 -> mkLHsPatTup [p1, p2]) pats
789 (rhs, _) = foldr1 (\(e1,t1) (e2,t2) ->
790 (mkApps mzip_op' [Type t1, Type t2, e1, e2],
791 mkBoxedTupleTy [t1,t2]))
792 exps_w_tys
793
794 ; dsMcBindStmt pat rhs bind_op noSyntaxExpr bind_ty stmts_rest }
795 where
796 ds_inner (ParStmtBlock _ stmts bndrs return_op)
797 = do { exp <- dsInnerMonadComp stmts bndrs return_op
798 ; return (exp, mkBigCoreVarTupTy bndrs) }
799 ds_inner (XParStmtBlock{}) = panic "dsMcStmt"
800
801 dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
802
803
804 matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
805 -- (matchTuple [a,b,c] body)
806 -- returns the Core term
807 -- \x. case x of (a,b,c) -> body
808 matchTuple ids body
809 = do { us <- newUniqueSupply
810 ; tup_id <- newSysLocalDs (mkBigCoreVarTupTy ids)
811 ; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
812
813 -- general `rhs' >>= \pat -> stmts` desugaring where `rhs'` is already a
814 -- desugared `CoreExpr`
815 dsMcBindStmt :: LPat GhcTc
816 -> CoreExpr -- ^ the desugared rhs of the bind statement
817 -> SyntaxExpr GhcTc
818 -> SyntaxExpr GhcTc
819 -> Type -- ^ S in (>>=) :: Q -> (R -> S) -> T
820 -> [ExprLStmt GhcTc]
821 -> DsM CoreExpr
822 dsMcBindStmt pat rhs' bind_op fail_op res1_ty stmts
823 = do { body <- dsMcStmts stmts
824 ; var <- selectSimpleMatchVarL pat
825 ; match <- matchSinglePat (Var var) (StmtCtxt DoExpr) pat
826 res1_ty (cantFailMatchResult body)
827 ; match_code <- handle_failure pat match fail_op
828 ; dsSyntaxExpr bind_op [rhs', Lam var match_code] }
829
830 where
831 -- In a monad comprehension expression, pattern-match failure just calls
832 -- the monadic `fail` rather than throwing an exception
833 handle_failure pat match fail_op
834 | matchCanFail match
835 = do { dflags <- getDynFlags
836 ; fail_msg <- mkStringExpr (mk_fail_msg dflags pat)
837 ; fail_expr <- dsSyntaxExpr fail_op [fail_msg]
838 ; extractMatchResult match fail_expr }
839 | otherwise
840 = extractMatchResult match (error "It can't fail")
841
842 mk_fail_msg :: DynFlags -> Located e -> String
843 mk_fail_msg dflags pat
844 = "Pattern match failure in monad comprehension at " ++
845 showPpr dflags (getLoc pat)
846
847 -- Desugar nested monad comprehensions, for example in `then..` constructs
848 -- dsInnerMonadComp quals [a,b,c] ret_op
849 -- returns the desugaring of
850 -- [ (a,b,c) | quals ]
851
852 dsInnerMonadComp :: [ExprLStmt GhcTc]
853 -> [Id] -- Return a tuple of these variables
854 -> SyntaxExpr GhcTc -- The monomorphic "return" operator
855 -> DsM CoreExpr
856 dsInnerMonadComp stmts bndrs ret_op
857 = dsMcStmts (stmts ++ [noLoc (LastStmt (mkBigLHsVarTupId bndrs) False ret_op)])
858
859
860 -- The `unzip` function for `GroupStmt` in a monad comprehensions
861 --
862 -- unzip :: m (a,b,..) -> (m a,m b,..)
863 -- unzip m_tuple = ( liftM selN1 m_tuple
864 -- , liftM selN2 m_tuple
865 -- , .. )
866 --
867 -- mkMcUnzipM fmap ys [t1, t2]
868 -- = ( fmap (selN1 :: (t1, t2) -> t1) ys
869 -- , fmap (selN2 :: (t1, t2) -> t2) ys )
870
871 mkMcUnzipM :: TransForm
872 -> HsExpr GhcTcId -- fmap
873 -> Id -- Of type n (a,b,c)
874 -> [Type] -- [a,b,c] (not levity-polymorphic)
875 -> DsM CoreExpr -- Of type (n a, n b, n c)
876 mkMcUnzipM ThenForm _ ys _
877 = return (Var ys) -- No unzipping to do
878
879 mkMcUnzipM _ fmap_op ys elt_tys
880 = do { fmap_op' <- dsExpr fmap_op
881 ; xs <- mapM newSysLocalDs elt_tys
882 ; let tup_ty = mkBigCoreTupTy elt_tys
883 ; tup_xs <- newSysLocalDs tup_ty
884
885 ; let mk_elt i = mkApps fmap_op' -- fmap :: forall a b. (a -> b) -> n a -> n b
886 [ Type tup_ty, Type (getNth elt_tys i)
887 , mk_sel i, Var ys]
888
889 mk_sel n = Lam tup_xs $
890 mkTupleSelector xs (getNth xs n) tup_xs (Var tup_xs)
891
892 ; return (mkBigCoreTup (map mk_elt [0..length elt_tys - 1])) }