Small refactor in desugar of pattern matching
[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, 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 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 deListComp (XStmtLR {} : _) _ =
270 panic "deListComp XStmtLR"
271
272 deBindComp :: OutPat GhcTc
273 -> CoreExpr
274 -> [ExprStmt GhcTc]
275 -> CoreExpr
276 -> DsM (Expr Id)
277 deBindComp pat core_list1 quals core_list2 = do
278 let u3_ty@u1_ty = exprType core_list1 -- two names, same thing
279
280 -- u1_ty is a [alpha] type, and u2_ty = alpha
281 let u2_ty = hsLPatType pat
282
283 let res_ty = exprType core_list2
284 h_ty = u1_ty `mkFunTy` res_ty
285
286 -- no levity polymorphism here, as list comprehensions don't work
287 -- with RebindableSyntax. NB: These are *not* monad comps.
288 [h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]
289
290 -- the "fail" value ...
291 let
292 core_fail = App (Var h) (Var u3)
293 letrec_body = App (Var h) core_list1
294
295 rest_expr <- deListComp quals core_fail
296 core_match <- matchSimply (Var u2) (StmtCtxt ListComp) pat rest_expr core_fail
297
298 let
299 rhs = Lam u1 $
300 Case (Var u1) u1 res_ty
301 [(DataAlt nilDataCon, [], core_list2),
302 (DataAlt consDataCon, [u2, u3], core_match)]
303 -- Increasing order of tag
304
305 return (Let (Rec [(h, rhs)]) letrec_body)
306
307 {-
308 ************************************************************************
309 * *
310 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
311 * *
312 ************************************************************************
313
314 @dfListComp@ are the rules used with foldr/build turned on:
315
316 \begin{verbatim}
317 TE[ e | ] c n = c e n
318 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
319 TE[ e | p <- l , q ] c n = let
320 f = \ x b -> case x of
321 p -> TE[ e | q ] c b
322 _ -> b
323 in
324 foldr f n l
325 \end{verbatim}
326 -}
327
328 dfListComp :: Id -> Id -- 'c' and 'n'
329 -> [ExprStmt GhcTc] -- the rest of the qual's
330 -> DsM CoreExpr
331
332 dfListComp _ _ [] = panic "dfListComp"
333
334 dfListComp c_id n_id (LastStmt _ body _ _ : quals)
335 = ASSERT( null quals )
336 do { core_body <- dsLExprNoLP body
337 ; return (mkApps (Var c_id) [core_body, Var n_id]) }
338
339 -- Non-last: must be a guard
340 dfListComp c_id n_id (BodyStmt _ guard _ _ : quals) = do
341 core_guard <- dsLExpr guard
342 core_rest <- dfListComp c_id n_id quals
343 return (mkIfThenElse core_guard core_rest (Var n_id))
344
345 dfListComp c_id n_id (LetStmt _ binds : quals) = do
346 -- new in 1.3, local bindings
347 core_rest <- dfListComp c_id n_id quals
348 dsLocalBinds binds core_rest
349
350 dfListComp c_id n_id (stmt@(TransStmt {}) : quals) = do
351 (inner_list_expr, pat) <- dsTransStmt stmt
352 -- Anyway, we bind the newly grouped list via the generic binding function
353 dfBindComp c_id n_id (pat, inner_list_expr) quals
354
355 dfListComp c_id n_id (BindStmt _ pat list1 _ _ : quals) = do
356 -- evaluate the two lists
357 core_list1 <- dsLExpr list1
358
359 -- Do the rest of the work in the generic binding builder
360 dfBindComp c_id n_id (pat, core_list1) quals
361
362 dfListComp _ _ (ParStmt {} : _) = panic "dfListComp ParStmt"
363 dfListComp _ _ (RecStmt {} : _) = panic "dfListComp RecStmt"
364 dfListComp _ _ (ApplicativeStmt {} : _) =
365 panic "dfListComp ApplicativeStmt"
366 dfListComp _ _ (XStmtLR {} : _) =
367 panic "dfListComp XStmtLR"
368
369 dfBindComp :: Id -> Id -- 'c' and 'n'
370 -> (LPat GhcTc, CoreExpr)
371 -> [ExprStmt GhcTc] -- the rest of the qual's
372 -> DsM CoreExpr
373 dfBindComp c_id n_id (pat, core_list1) quals = do
374 -- find the required type
375 let x_ty = hsLPatType pat
376 let b_ty = idType n_id
377
378 -- create some new local id's
379 b <- newSysLocalDs b_ty
380 x <- newSysLocalDs x_ty
381
382 -- build rest of the comprehesion
383 core_rest <- dfListComp c_id b quals
384
385 -- build the pattern match
386 core_expr <- matchSimply (Var x) (StmtCtxt ListComp)
387 pat core_rest (Var b)
388
389 -- now build the outermost foldr, and return
390 mkFoldrExpr x_ty b_ty (mkLams [x, b] core_expr) (Var n_id) core_list1
391
392 {-
393 ************************************************************************
394 * *
395 \subsection[DsFunGeneration]{Generation of zip/unzip functions for use in desugaring}
396 * *
397 ************************************************************************
398 -}
399
400 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
401 -- mkZipBind [t1, t2]
402 -- = (zip, \as1:[t1] as2:[t2]
403 -- -> case as1 of
404 -- [] -> []
405 -- (a1:as'1) -> case as2 of
406 -- [] -> []
407 -- (a2:as'2) -> (a1, a2) : zip as'1 as'2)]
408
409 mkZipBind elt_tys = do
410 ass <- mapM newSysLocalDs elt_list_tys
411 as' <- mapM newSysLocalDs elt_tys
412 as's <- mapM newSysLocalDs elt_list_tys
413
414 zip_fn <- newSysLocalDs zip_fn_ty
415
416 let inner_rhs = mkConsExpr elt_tuple_ty
417 (mkBigCoreVarTup as')
418 (mkVarApps (Var zip_fn) as's)
419 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
420
421 return (zip_fn, mkLams ass zip_body)
422 where
423 elt_list_tys = map mkListTy elt_tys
424 elt_tuple_ty = mkBigCoreTupTy elt_tys
425 elt_tuple_list_ty = mkListTy elt_tuple_ty
426
427 zip_fn_ty = mkFunTys elt_list_tys elt_tuple_list_ty
428
429 mk_case (as, a', as') rest
430 = Case (Var as) as elt_tuple_list_ty
431 [(DataAlt nilDataCon, [], mkNilExpr elt_tuple_ty),
432 (DataAlt consDataCon, [a', as'], rest)]
433 -- Increasing order of tag
434
435
436 mkUnzipBind :: TransForm -> [Type] -> DsM (Maybe (Id, CoreExpr))
437 -- mkUnzipBind [t1, t2]
438 -- = (unzip, \ys :: [(t1, t2)] -> foldr (\ax :: (t1, t2) axs :: ([t1], [t2])
439 -- -> case ax of
440 -- (x1, x2) -> case axs of
441 -- (xs1, xs2) -> (x1 : xs1, x2 : xs2))
442 -- ([], [])
443 -- ys)
444 --
445 -- We use foldr here in all cases, even if rules are turned off, because we may as well!
446 mkUnzipBind ThenForm _
447 = return Nothing -- No unzipping for ThenForm
448 mkUnzipBind _ elt_tys
449 = do { ax <- newSysLocalDs elt_tuple_ty
450 ; axs <- newSysLocalDs elt_list_tuple_ty
451 ; ys <- newSysLocalDs elt_tuple_list_ty
452 ; xs <- mapM newSysLocalDs elt_tys
453 ; xss <- mapM newSysLocalDs elt_list_tys
454
455 ; unzip_fn <- newSysLocalDs unzip_fn_ty
456
457 ; [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
458
459 ; let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
460 concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
461 tupled_concat_expression = mkBigCoreTup concat_expressions
462
463 folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
464 folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
465 folder_body = mkLams [ax, axs] folder_body_outer_case
466
467 ; unzip_body <- mkFoldrExpr elt_tuple_ty elt_list_tuple_ty folder_body nil_tuple (Var ys)
468 ; return (Just (unzip_fn, mkLams [ys] unzip_body)) }
469 where
470 elt_tuple_ty = mkBigCoreTupTy elt_tys
471 elt_tuple_list_ty = mkListTy elt_tuple_ty
472 elt_list_tys = map mkListTy elt_tys
473 elt_list_tuple_ty = mkBigCoreTupTy elt_list_tys
474
475 unzip_fn_ty = elt_tuple_list_ty `mkFunTy` elt_list_tuple_ty
476
477 mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
478
479 -- Translation for monad comprehensions
480
481 -- Entry point for monad comprehension desugaring
482 dsMonadComp :: [ExprLStmt GhcTc] -> DsM CoreExpr
483 dsMonadComp stmts = dsMcStmts stmts
484
485 dsMcStmts :: [ExprLStmt GhcTc] -> DsM CoreExpr
486 dsMcStmts [] = panic "dsMcStmts"
487 dsMcStmts (L loc stmt : lstmts) = putSrcSpanDs loc (dsMcStmt stmt lstmts)
488
489 ---------------
490 dsMcStmt :: ExprStmt GhcTc -> [ExprLStmt GhcTc] -> DsM CoreExpr
491
492 dsMcStmt (LastStmt _ body _ ret_op) stmts
493 = ASSERT( null stmts )
494 do { body' <- dsLExpr body
495 ; dsSyntaxExpr ret_op [body'] }
496
497 -- [ .. | let binds, stmts ]
498 dsMcStmt (LetStmt _ binds) stmts
499 = do { rest <- dsMcStmts stmts
500 ; dsLocalBinds binds rest }
501
502 -- [ .. | a <- m, stmts ]
503 dsMcStmt (BindStmt bind_ty pat rhs bind_op fail_op) stmts
504 = do { rhs' <- dsLExpr rhs
505 ; dsMcBindStmt pat rhs' bind_op fail_op bind_ty stmts }
506
507 -- Apply `guard` to the `exp` expression
508 --
509 -- [ .. | exp, stmts ]
510 --
511 dsMcStmt (BodyStmt _ exp then_exp guard_exp) stmts
512 = do { exp' <- dsLExpr exp
513 ; rest <- dsMcStmts stmts
514 ; guard_exp' <- dsSyntaxExpr guard_exp [exp']
515 ; dsSyntaxExpr then_exp [guard_exp', rest] }
516
517 -- Group statements desugar like this:
518 --
519 -- [| (q, then group by e using f); rest |]
520 -- ---> f {qt} (\qv -> e) [| q; return qv |] >>= \ n_tup ->
521 -- case unzip n_tup of qv' -> [| rest |]
522 --
523 -- where variables (v1:t1, ..., vk:tk) are bound by q
524 -- qv = (v1, ..., vk)
525 -- qt = (t1, ..., tk)
526 -- (>>=) :: m2 a -> (a -> m3 b) -> m3 b
527 -- f :: forall a. (a -> t) -> m1 a -> m2 (n a)
528 -- n_tup :: n qt
529 -- unzip :: n qt -> (n t1, ..., n tk) (needs Functor n)
530
531 dsMcStmt (TransStmt { trS_stmts = stmts, trS_bndrs = bndrs
532 , trS_by = by, trS_using = using
533 , trS_ret = return_op, trS_bind = bind_op
534 , trS_ext = n_tup_ty' -- n (a,b,c)
535 , trS_fmap = fmap_op, trS_form = form }) stmts_rest
536 = do { let (from_bndrs, to_bndrs) = unzip bndrs
537
538 ; let from_bndr_tys = map idType from_bndrs -- Types ty
539
540
541 -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
542 ; expr' <- dsInnerMonadComp stmts from_bndrs return_op
543
544 -- Work out what arguments should be supplied to that expression: i.e. is an extraction
545 -- function required? If so, create that desugared function and add to arguments
546 ; usingExpr' <- dsLExpr using
547 ; usingArgs' <- case by of
548 Nothing -> return [expr']
549 Just by_e -> do { by_e' <- dsLExpr by_e
550 ; lam' <- matchTuple from_bndrs by_e'
551 ; return [lam', expr'] }
552
553 -- Generate the expressions to build the grouped list
554 -- Build a pattern that ensures the consumer binds into the NEW binders,
555 -- which hold monads rather than single values
556 ; let tup_n_ty' = mkBigCoreVarTupTy to_bndrs
557
558 ; body <- dsMcStmts stmts_rest
559 ; n_tup_var' <- newSysLocalDsNoLP n_tup_ty'
560 ; tup_n_var' <- newSysLocalDs tup_n_ty'
561 ; tup_n_expr' <- mkMcUnzipM form fmap_op n_tup_var' from_bndr_tys
562 ; us <- newUniqueSupply
563 ; let rhs' = mkApps usingExpr' usingArgs'
564 body' = mkTupleCase us to_bndrs body tup_n_var' tup_n_expr'
565
566 ; dsSyntaxExpr bind_op [rhs', Lam n_tup_var' body'] }
567
568 -- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel
569 -- statements, for example:
570 --
571 -- [ body | qs1 | qs2 | qs3 ]
572 -- -> [ body | (bndrs1, (bndrs2, bndrs3))
573 -- <- [bndrs1 | qs1] `mzip` ([bndrs2 | qs2] `mzip` [bndrs3 | qs3]) ]
574 --
575 -- where `mzip` has type
576 -- mzip :: forall a b. m a -> m b -> m (a,b)
577 -- NB: we need a polymorphic mzip because we call it several times
578
579 dsMcStmt (ParStmt bind_ty blocks mzip_op bind_op) stmts_rest
580 = do { exps_w_tys <- mapM ds_inner blocks -- Pairs (exp :: m ty, ty)
581 ; mzip_op' <- dsExpr mzip_op
582
583 ; let -- The pattern variables
584 pats = [ mkBigLHsVarPatTupId bs | ParStmtBlock _ _ bs _ <- blocks]
585 -- Pattern with tuples of variables
586 -- [v1,v2,v3] => (v1, (v2, v3))
587 pat = foldr1 (\p1 p2 -> mkLHsPatTup [p1, p2]) pats
588 (rhs, _) = foldr1 (\(e1,t1) (e2,t2) ->
589 (mkApps mzip_op' [Type t1, Type t2, e1, e2],
590 mkBoxedTupleTy [t1,t2]))
591 exps_w_tys
592
593 ; dsMcBindStmt pat rhs bind_op noSyntaxExpr bind_ty stmts_rest }
594 where
595 ds_inner (ParStmtBlock _ stmts bndrs return_op)
596 = do { exp <- dsInnerMonadComp stmts bndrs return_op
597 ; return (exp, mkBigCoreVarTupTy bndrs) }
598 ds_inner (XParStmtBlock{}) = panic "dsMcStmt"
599
600 dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
601
602
603 matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
604 -- (matchTuple [a,b,c] body)
605 -- returns the Core term
606 -- \x. case x of (a,b,c) -> body
607 matchTuple ids body
608 = do { us <- newUniqueSupply
609 ; tup_id <- newSysLocalDs (mkBigCoreVarTupTy ids)
610 ; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
611
612 -- general `rhs' >>= \pat -> stmts` desugaring where `rhs'` is already a
613 -- desugared `CoreExpr`
614 dsMcBindStmt :: LPat GhcTc
615 -> CoreExpr -- ^ the desugared rhs of the bind statement
616 -> SyntaxExpr GhcTc
617 -> SyntaxExpr GhcTc
618 -> Type -- ^ S in (>>=) :: Q -> (R -> S) -> T
619 -> [ExprLStmt GhcTc]
620 -> DsM CoreExpr
621 dsMcBindStmt pat rhs' bind_op fail_op res1_ty stmts
622 = do { body <- dsMcStmts stmts
623 ; var <- selectSimpleMatchVarL pat
624 ; match <- matchSinglePatVar var (StmtCtxt DoExpr) pat
625 res1_ty (cantFailMatchResult body)
626 ; match_code <- handle_failure pat match fail_op
627 ; dsSyntaxExpr bind_op [rhs', Lam var match_code] }
628
629 where
630 -- In a monad comprehension expression, pattern-match failure just calls
631 -- the monadic `fail` rather than throwing an exception
632 handle_failure pat match fail_op
633 | matchCanFail match
634 = do { dflags <- getDynFlags
635 ; fail_msg <- mkStringExpr (mk_fail_msg dflags pat)
636 ; fail_expr <- dsSyntaxExpr fail_op [fail_msg]
637 ; extractMatchResult match fail_expr }
638 | otherwise
639 = extractMatchResult match (error "It can't fail")
640
641 mk_fail_msg :: DynFlags -> Located e -> String
642 mk_fail_msg dflags pat
643 = "Pattern match failure in monad comprehension at " ++
644 showPpr dflags (getLoc pat)
645
646 -- Desugar nested monad comprehensions, for example in `then..` constructs
647 -- dsInnerMonadComp quals [a,b,c] ret_op
648 -- returns the desugaring of
649 -- [ (a,b,c) | quals ]
650
651 dsInnerMonadComp :: [ExprLStmt GhcTc]
652 -> [Id] -- Return a tuple of these variables
653 -> SyntaxExpr GhcTc -- The monomorphic "return" operator
654 -> DsM CoreExpr
655 dsInnerMonadComp stmts bndrs ret_op
656 = dsMcStmts (stmts ++
657 [noLoc (LastStmt noExt (mkBigLHsVarTupId bndrs) False ret_op)])
658
659
660 -- The `unzip` function for `GroupStmt` in a monad comprehensions
661 --
662 -- unzip :: m (a,b,..) -> (m a,m b,..)
663 -- unzip m_tuple = ( liftM selN1 m_tuple
664 -- , liftM selN2 m_tuple
665 -- , .. )
666 --
667 -- mkMcUnzipM fmap ys [t1, t2]
668 -- = ( fmap (selN1 :: (t1, t2) -> t1) ys
669 -- , fmap (selN2 :: (t1, t2) -> t2) ys )
670
671 mkMcUnzipM :: TransForm
672 -> HsExpr GhcTcId -- fmap
673 -> Id -- Of type n (a,b,c)
674 -> [Type] -- [a,b,c] (not levity-polymorphic)
675 -> DsM CoreExpr -- Of type (n a, n b, n c)
676 mkMcUnzipM ThenForm _ ys _
677 = return (Var ys) -- No unzipping to do
678
679 mkMcUnzipM _ fmap_op ys elt_tys
680 = do { fmap_op' <- dsExpr fmap_op
681 ; xs <- mapM newSysLocalDs elt_tys
682 ; let tup_ty = mkBigCoreTupTy elt_tys
683 ; tup_xs <- newSysLocalDs tup_ty
684
685 ; let mk_elt i = mkApps fmap_op' -- fmap :: forall a b. (a -> b) -> n a -> n b
686 [ Type tup_ty, Type (getNth elt_tys i)
687 , mk_sel i, Var ys]
688
689 mk_sel n = Lam tup_xs $
690 mkTupleSelector xs (getNth xs n) tup_xs (Var tup_xs)
691
692 ; return (mkBigCoreTup (map mk_elt [0..length elt_tys - 1])) }