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