never executed always true always false
1
2 {-# LANGUAGE TypeFamilies #-}
3
4 {-
5 (c) The University of Glasgow 2006
6 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
7
8
9 Desugaring list comprehensions, monad comprehensions and array comprehensions
10 -}
11
12 module GHC.HsToCore.ListComp ( dsListComp, dsMonadComp ) where
13
14 import GHC.Prelude
15
16 import {-# SOURCE #-} GHC.HsToCore.Expr ( dsExpr, dsLExpr, dsLocalBinds, dsSyntaxExpr )
17
18 import GHC.Hs
19 import GHC.Hs.Syn.Type
20 import GHC.Core
21 import GHC.Core.Make
22
23 import GHC.HsToCore.Monad -- the monadery used in the desugarer
24 import GHC.HsToCore.Utils
25
26 import GHC.Driver.Session
27 import GHC.Core.Utils
28 import GHC.Types.Id
29 import GHC.Core.Type
30 import GHC.Builtin.Types
31 import GHC.HsToCore.Match
32 import GHC.Builtin.Names
33 import GHC.Types.SrcLoc
34 import GHC.Utils.Outputable
35 import GHC.Utils.Panic
36 import GHC.Utils.Panic.Plain
37 import GHC.Tc.Utils.TcType
38 import GHC.Data.List.SetOps( getNth )
39
40 {-
41 List comprehensions may be desugared in one of two ways: ``ordinary''
42 (as you would expect if you read SLPJ's book) and ``with foldr/build
43 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
44
45 There will be at least one ``qualifier'' in the input.
46 -}
47
48 dsListComp :: [ExprLStmt GhcTc]
49 -> Type -- Type of entire list
50 -> DsM CoreExpr
51 dsListComp lquals res_ty = do
52 dflags <- getDynFlags
53 let quals = map unLoc lquals
54 elt_ty = case tcTyConAppArgs res_ty of
55 [elt_ty] -> elt_ty
56 _ -> pprPanic "dsListComp" (ppr res_ty $$ ppr lquals)
57
58 if not (gopt Opt_EnableRewriteRules dflags) || gopt Opt_IgnoreInterfacePragmas dflags
59 -- Either rules are switched off, or we are ignoring what there are;
60 -- Either way foldr/build won't happen, so use the more efficient
61 -- Wadler-style desugaring
62 || isParallelComp quals
63 -- Foldr-style desugaring can't handle parallel list comprehensions
64 then deListComp quals (mkNilExpr elt_ty)
65 else mkBuildExpr elt_ty (\(c, _) (n, _) -> dfListComp c n quals)
66 -- Foldr/build should be enabled, so desugar
67 -- into foldrs and builds
68
69 where
70 -- We must test for ParStmt anywhere, not just at the head, because an extension
71 -- to list comprehensions would be to add brackets to specify the associativity
72 -- of qualifier lists. This is really easy to do by adding extra ParStmts into the
73 -- mix of possibly a single element in length, so we do this to leave the possibility open
74 isParallelComp = any isParallelStmt
75
76 isParallelStmt (ParStmt {}) = True
77 isParallelStmt _ = False
78
79
80 -- This function lets you desugar a inner list comprehension and a list of the binders
81 -- of that comprehension that we need in the outer comprehension into such an expression
82 -- and the type of the elements that it outputs (tuples of binders)
83 dsInnerListComp :: (ParStmtBlock GhcTc GhcTc) -> DsM (CoreExpr, Type)
84 dsInnerListComp (ParStmtBlock _ stmts bndrs _)
85 = do { let bndrs_tuple_type = mkBigCoreVarTupTy bndrs
86 list_ty = mkListTy bndrs_tuple_type
87
88 -- really use original bndrs below!
89 ; expr <- dsListComp (stmts ++ [noLocA $ mkLastStmt (mkBigLHsVarTupId bndrs)]) list_ty
90
91 ; return (expr, bndrs_tuple_type) }
92
93 -- This function factors out commonality between the desugaring strategies for GroupStmt.
94 -- Given such a statement it gives you back an expression representing how to compute the transformed
95 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
96 dsTransStmt :: ExprStmt GhcTc -> DsM (CoreExpr, LPat GhcTc)
97 dsTransStmt (TransStmt { trS_form = form, trS_stmts = stmts, trS_bndrs = binderMap
98 , trS_by = by, trS_using = using }) = do
99 let (from_bndrs, to_bndrs) = unzip binderMap
100
101 let from_bndrs_tys = map idType from_bndrs
102 to_bndrs_tys = map idType to_bndrs
103
104 to_bndrs_tup_ty = mkBigCoreTupTy to_bndrs_tys
105
106 -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
107 (expr', from_tup_ty) <- dsInnerListComp (ParStmtBlock noExtField stmts
108 from_bndrs noSyntaxExpr)
109
110 -- Work out what arguments should be supplied to that expression: i.e. is an extraction
111 -- function required? If so, create that desugared function and add to arguments
112 usingExpr' <- dsLExpr using
113 usingArgs' <- case by of
114 Nothing -> return [expr']
115 Just by_e -> do { by_e' <- dsLExpr by_e
116 ; lam' <- matchTuple from_bndrs by_e'
117 ; return [lam', expr'] }
118
119 -- Create an unzip function for the appropriate arity and element types and find "map"
120 unzip_stuff' <- mkUnzipBind form from_bndrs_tys
121 map_id <- dsLookupGlobalId mapName
122
123 -- Generate the expressions to build the grouped list
124 let -- First we apply the grouping function to the inner list
125 inner_list_expr' = mkApps usingExpr' usingArgs'
126 -- Then we map our "unzip" across it to turn the lists of tuples into tuples of lists
127 -- We make sure we instantiate the type variable "a" to be a list of "from" tuples and
128 -- the "b" to be a tuple of "to" lists!
129 -- Then finally we bind the unzip function around that expression
130 bound_unzipped_inner_list_expr'
131 = case unzip_stuff' of
132 Nothing -> inner_list_expr'
133 Just (unzip_fn', unzip_rhs') ->
134 Let (Rec [(unzip_fn', unzip_rhs')]) $
135 mkApps (Var map_id) $
136 [ Type (mkListTy from_tup_ty)
137 , Type to_bndrs_tup_ty
138 , Var unzip_fn'
139 , inner_list_expr' ]
140
141 -- Build a pattern that ensures the consumer binds into the NEW binders,
142 -- which hold lists rather than single values
143 let pat = mkBigLHsVarPatTupId to_bndrs -- NB: no '!
144 return (bound_unzipped_inner_list_expr', pat)
145
146 dsTransStmt _ = panic "dsTransStmt: Not given a TransStmt"
147
148 {-
149 ************************************************************************
150 * *
151 * Ordinary desugaring of list comprehensions *
152 * *
153 ************************************************************************
154
155 Just as in Phil's chapter~7 in SLPJ, using the rules for
156 optimally-compiled list comprehensions. This is what Kevin followed
157 as well, and I quite happily do the same. The TQ translation scheme
158 transforms a list of qualifiers (either boolean expressions or
159 generators) into a single expression which implements the list
160 comprehension. Because we are generating 2nd-order polymorphic
161 lambda-calculus, calls to NIL and CONS must be applied to a type
162 argument, as well as their usual value arguments.
163 \begin{verbatim}
164 TE << [ e | qs ] >> = TQ << [ e | qs ] ++ Nil (typeOf e) >>
165
166 (Rule C)
167 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
168
169 (Rule B)
170 TQ << [ e | b , qs ] ++ L >> =
171 if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
172
173 (Rule A')
174 TQ << [ e | p <- L1, qs ] ++ L2 >> =
175 letrec
176 h = \ u1 ->
177 case u1 of
178 [] -> TE << L2 >>
179 (u2 : u3) ->
180 (( \ TE << p >> -> ( TQ << [e | qs] ++ (h u3) >> )) u2)
181 [] (h u3)
182 in
183 h ( TE << L1 >> )
184
185 "h", "u1", "u2", and "u3" are new variables.
186 \end{verbatim}
187
188 @deListComp@ is the TQ translation scheme. Roughly speaking, @dsExpr@
189 is the TE translation scheme. Note that we carry around the @L@ list
190 already desugared. @dsListComp@ does the top TE rule mentioned above.
191
192 To the above, we add an additional rule to deal with parallel list
193 comprehensions. The translation goes roughly as follows:
194 [ e | p1 <- e11, let v1 = e12, p2 <- e13
195 | q1 <- e21, let v2 = e22, q2 <- e23]
196 =>
197 [ e | ((x1, .., xn), (y1, ..., ym)) <-
198 zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
199 [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
200 where (x1, .., xn) are the variables bound in p1, v1, p2
201 (y1, .., ym) are the variables bound in q1, v2, q2
202
203 In the translation below, the ParStmt branch translates each parallel branch
204 into a sub-comprehension, and desugars each independently. The resulting lists
205 are fed to a zip function, we create a binding for all the variables bound in all
206 the comprehensions, and then we hand things off the desugarer for bindings.
207 The zip function is generated here a) because it's small, and b) because then we
208 don't have to deal with arbitrary limits on the number of zip functions in the
209 prelude, nor which library the zip function came from.
210 The introduced tuples are Boxed, but only because I couldn't get it to work
211 with the Unboxed variety.
212 -}
213
214 deListComp :: [ExprStmt GhcTc] -> CoreExpr -> DsM CoreExpr
215
216 deListComp [] _ = panic "deListComp"
217
218 deListComp (LastStmt _ body _ _ : quals) list
219 = -- Figure 7.4, SLPJ, p 135, rule C above
220 assert (null quals) $
221 do { core_body <- dsLExpr body
222 ; return (mkConsExpr (exprType core_body) core_body list) }
223
224 -- Non-last: must be a guard
225 deListComp (BodyStmt _ guard _ _ : quals) list = do -- rule B above
226 core_guard <- dsLExpr guard
227 core_rest <- deListComp quals list
228 return (mkIfThenElse core_guard core_rest list)
229
230 -- [e | let B, qs] = let B in [e | qs]
231 deListComp (LetStmt _ binds : quals) list = do
232 core_rest <- deListComp quals list
233 dsLocalBinds binds core_rest
234
235 deListComp (stmt@(TransStmt {}) : quals) list = do
236 (inner_list_expr, pat) <- dsTransStmt stmt
237 deBindComp pat inner_list_expr quals list
238
239 deListComp (BindStmt _ pat list1 : quals) core_list2 = do -- rule A' above
240 core_list1 <- dsLExpr list1
241 deBindComp pat core_list1 quals core_list2
242
243 deListComp (ParStmt _ stmtss_w_bndrs _ _ : quals) list
244 = do { exps_and_qual_tys <- mapM dsInnerListComp stmtss_w_bndrs
245 ; let (exps, qual_tys) = unzip exps_and_qual_tys
246
247 ; (zip_fn, zip_rhs) <- mkZipBind qual_tys
248
249 -- Deal with [e | pat <- zip l1 .. ln] in example above
250 ; deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
251 quals list }
252 where
253 bndrs_s = [bs | ParStmtBlock _ _ bs _ <- stmtss_w_bndrs]
254
255 -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
256 pat = mkBigLHsPatTupId pats
257 pats = map mkBigLHsVarPatTupId bndrs_s
258
259 deListComp (RecStmt {} : _) _ = panic "deListComp RecStmt"
260
261 deListComp (ApplicativeStmt {} : _) _ =
262 panic "deListComp ApplicativeStmt"
263
264 deBindComp :: LPat GhcTc
265 -> CoreExpr
266 -> [ExprStmt GhcTc]
267 -> CoreExpr
268 -> DsM (Expr Id)
269 deBindComp pat core_list1 quals core_list2 = do
270 let u3_ty@u1_ty = exprType core_list1 -- two names, same thing
271
272 -- u1_ty is a [alpha] type, and u2_ty = alpha
273 let u2_ty = hsLPatType pat
274
275 let res_ty = exprType core_list2
276 h_ty = u1_ty `mkVisFunTyMany` res_ty
277
278 -- no representation polymorphism here, as list comprehensions
279 -- don't work with RebindableSyntax. NB: These are *not* monad comps.
280 [h, u1, u2, u3] <- newSysLocalsDs $ map unrestricted [h_ty, u1_ty, u2_ty, u3_ty]
281
282 -- the "fail" value ...
283 let
284 core_fail = App (Var h) (Var u3)
285 letrec_body = App (Var h) core_list1
286
287 rest_expr <- deListComp quals core_fail
288 core_match <- matchSimply (Var u2) (StmtCtxt (HsDoStmt ListComp)) pat rest_expr core_fail
289
290 let
291 rhs = Lam u1 $
292 Case (Var u1) u1 res_ty
293 [Alt (DataAlt nilDataCon) [] core_list2
294 ,Alt (DataAlt consDataCon) [u2, u3] core_match]
295 -- Increasing order of tag
296
297 return (Let (Rec [(h, rhs)]) letrec_body)
298
299 {-
300 ************************************************************************
301 * *
302 * Foldr/Build desugaring of list comprehensions *
303 * *
304 ************************************************************************
305
306 @dfListComp@ are the rules used with foldr/build turned on:
307
308 \begin{verbatim}
309 TE[ e | ] c n = c e n
310 TE[ e | b , q ] c n = if b then TE[ e | q ] c n else n
311 TE[ e | p <- l , q ] c n = let
312 f = \ x b -> case x of
313 p -> TE[ e | q ] c b
314 _ -> b
315 in
316 foldr f n l
317 \end{verbatim}
318 -}
319
320 dfListComp :: Id -> Id -- 'c' and 'n'
321 -> [ExprStmt GhcTc] -- the rest of the qual's
322 -> DsM CoreExpr
323
324 dfListComp _ _ [] = panic "dfListComp"
325
326 dfListComp c_id n_id (LastStmt _ body _ _ : quals)
327 = assert (null quals) $
328 do { core_body <- dsLExpr body
329 ; return (mkApps (Var c_id) [core_body, Var n_id]) }
330
331 -- Non-last: must be a guard
332 dfListComp c_id n_id (BodyStmt _ guard _ _ : quals) = do
333 core_guard <- dsLExpr guard
334 core_rest <- dfListComp c_id n_id quals
335 return (mkIfThenElse core_guard core_rest (Var n_id))
336
337 dfListComp c_id n_id (LetStmt _ binds : quals) = do
338 -- new in 1.3, local bindings
339 core_rest <- dfListComp c_id n_id quals
340 dsLocalBinds binds core_rest
341
342 dfListComp c_id n_id (stmt@(TransStmt {}) : quals) = do
343 (inner_list_expr, pat) <- dsTransStmt stmt
344 -- Anyway, we bind the newly grouped list via the generic binding function
345 dfBindComp c_id n_id (pat, inner_list_expr) quals
346
347 dfListComp c_id n_id (BindStmt _ pat list1 : quals) = do
348 -- evaluate the two lists
349 core_list1 <- dsLExpr list1
350
351 -- Do the rest of the work in the generic binding builder
352 dfBindComp c_id n_id (pat, core_list1) quals
353
354 dfListComp _ _ (ParStmt {} : _) = panic "dfListComp ParStmt"
355 dfListComp _ _ (RecStmt {} : _) = panic "dfListComp RecStmt"
356 dfListComp _ _ (ApplicativeStmt {} : _) =
357 panic "dfListComp ApplicativeStmt"
358
359 dfBindComp :: Id -> Id -- 'c' and 'n'
360 -> (LPat GhcTc, CoreExpr)
361 -> [ExprStmt GhcTc] -- the rest of the qual's
362 -> DsM CoreExpr
363 dfBindComp c_id n_id (pat, core_list1) quals = do
364 -- find the required type
365 let x_ty = hsLPatType pat
366 let b_ty = idType n_id
367
368 -- create some new local id's
369 b <- newSysLocalDs Many b_ty
370 x <- newSysLocalDs Many x_ty
371
372 -- build rest of the comprehension
373 core_rest <- dfListComp c_id b quals
374
375 -- build the pattern match
376 core_expr <- matchSimply (Var x) (StmtCtxt (HsDoStmt ListComp))
377 pat core_rest (Var b)
378
379 -- now build the outermost foldr, and return
380 mkFoldrExpr x_ty b_ty (mkLams [x, b] core_expr) (Var n_id) core_list1
381
382 {-
383 ************************************************************************
384 * *
385 \subsection[DsFunGeneration]{Generation of zip/unzip functions for use in desugaring}
386 * *
387 ************************************************************************
388 -}
389
390 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
391 -- mkZipBind [t1, t2]
392 -- = (zip, \as1:[t1] as2:[t2]
393 -- -> case as1 of
394 -- [] -> []
395 -- (a1:as'1) -> case as2 of
396 -- [] -> []
397 -- (a2:as'2) -> (a1, a2) : zip as'1 as'2)]
398
399 mkZipBind elt_tys = do
400 ass <- mapM (newSysLocalDs Many) elt_list_tys
401 as' <- mapM (newSysLocalDs Many) elt_tys
402 as's <- mapM (newSysLocalDs Many) elt_list_tys
403
404 zip_fn <- newSysLocalDs Many zip_fn_ty
405
406 let inner_rhs = mkConsExpr elt_tuple_ty
407 (mkBigCoreVarTup as')
408 (mkVarApps (Var zip_fn) as's)
409 zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
410
411 return (zip_fn, mkLams ass zip_body)
412 where
413 elt_list_tys = map mkListTy elt_tys
414 elt_tuple_ty = mkBigCoreTupTy elt_tys
415 elt_tuple_list_ty = mkListTy elt_tuple_ty
416
417 zip_fn_ty = mkVisFunTysMany elt_list_tys elt_tuple_list_ty
418
419 mk_case (as, a', as') rest
420 = Case (Var as) as elt_tuple_list_ty
421 [ Alt (DataAlt nilDataCon) [] (mkNilExpr elt_tuple_ty)
422 , Alt (DataAlt consDataCon) [a', as'] rest]
423 -- Increasing order of tag
424
425
426 mkUnzipBind :: TransForm -> [Type] -> DsM (Maybe (Id, CoreExpr))
427 -- mkUnzipBind [t1, t2]
428 -- = (unzip, \ys :: [(t1, t2)] -> foldr (\ax :: (t1, t2) axs :: ([t1], [t2])
429 -- -> case ax of
430 -- (x1, x2) -> case axs of
431 -- (xs1, xs2) -> (x1 : xs1, x2 : xs2))
432 -- ([], [])
433 -- ys)
434 --
435 -- We use foldr here in all cases, even if rules are turned off, because we may as well!
436 mkUnzipBind ThenForm _
437 = return Nothing -- No unzipping for ThenForm
438 mkUnzipBind _ elt_tys
439 = do { ax <- newSysLocalDs Many elt_tuple_ty
440 ; axs <- newSysLocalDs Many elt_list_tuple_ty
441 ; ys <- newSysLocalDs Many elt_tuple_list_ty
442 ; xs <- mapM (newSysLocalDs Many) elt_tys
443 ; xss <- mapM (newSysLocalDs Many) elt_list_tys
444
445 ; unzip_fn <- newSysLocalDs Many unzip_fn_ty
446
447 ; [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
448
449 ; let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
450 concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
451 tupled_concat_expression = mkBigCoreTup concat_expressions
452
453 folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
454 folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
455 folder_body = mkLams [ax, axs] folder_body_outer_case
456
457 ; unzip_body <- mkFoldrExpr elt_tuple_ty elt_list_tuple_ty folder_body nil_tuple (Var ys)
458 ; return (Just (unzip_fn, mkLams [ys] unzip_body)) }
459 where
460 elt_tuple_ty = mkBigCoreTupTy elt_tys
461 elt_tuple_list_ty = mkListTy elt_tuple_ty
462 elt_list_tys = map mkListTy elt_tys
463 elt_list_tuple_ty = mkBigCoreTupTy elt_list_tys
464
465 unzip_fn_ty = elt_tuple_list_ty `mkVisFunTyMany` elt_list_tuple_ty
466
467 mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
468
469 -- Translation for monad comprehensions
470
471 -- Entry point for monad comprehension desugaring
472 dsMonadComp :: [ExprLStmt GhcTc] -> DsM CoreExpr
473 dsMonadComp stmts = dsMcStmts stmts
474
475 dsMcStmts :: [ExprLStmt GhcTc] -> DsM CoreExpr
476 dsMcStmts [] = panic "dsMcStmts"
477 dsMcStmts ((L loc stmt) : lstmts) = putSrcSpanDsA loc (dsMcStmt stmt lstmts)
478
479 ---------------
480 dsMcStmt :: ExprStmt GhcTc -> [ExprLStmt GhcTc] -> DsM CoreExpr
481
482 dsMcStmt (LastStmt _ body _ ret_op) stmts
483 = assert (null stmts) $
484 do { body' <- dsLExpr body
485 ; dsSyntaxExpr ret_op [body'] }
486
487 -- [ .. | let binds, stmts ]
488 dsMcStmt (LetStmt _ binds) stmts
489 = do { rest <- dsMcStmts stmts
490 ; dsLocalBinds binds rest }
491
492 -- [ .. | a <- m, stmts ]
493 dsMcStmt (BindStmt xbs pat rhs) stmts
494 = do { rhs' <- dsLExpr rhs
495 ; dsMcBindStmt pat rhs' (xbstc_bindOp xbs) (xbstc_failOp xbs) (xbstc_boundResultType xbs) stmts }
496
497 -- Apply `guard` to the `exp` expression
498 --
499 -- [ .. | exp, stmts ]
500 --
501 dsMcStmt (BodyStmt _ exp then_exp guard_exp) stmts
502 = do { exp' <- dsLExpr exp
503 ; rest <- dsMcStmts stmts
504 ; guard_exp' <- dsSyntaxExpr guard_exp [exp']
505 ; dsSyntaxExpr then_exp [guard_exp', rest] }
506
507 -- Group statements desugar like this:
508 --
509 -- [| (q, then group by e using f); rest |]
510 -- ---> f {qt} (\qv -> e) [| q; return qv |] >>= \ n_tup ->
511 -- case unzip n_tup of qv' -> [| rest |]
512 --
513 -- where variables (v1:t1, ..., vk:tk) are bound by q
514 -- qv = (v1, ..., vk)
515 -- qt = (t1, ..., tk)
516 -- (>>=) :: m2 a -> (a -> m3 b) -> m3 b
517 -- f :: forall a. (a -> t) -> m1 a -> m2 (n a)
518 -- n_tup :: n qt
519 -- unzip :: n qt -> (n t1, ..., n tk) (needs Functor n)
520
521 dsMcStmt (TransStmt { trS_stmts = stmts, trS_bndrs = bndrs
522 , trS_by = by, trS_using = using
523 , trS_ret = return_op, trS_bind = bind_op
524 , trS_ext = n_tup_ty' -- n (a,b,c)
525 , trS_fmap = fmap_op, trS_form = form }) stmts_rest
526 = do { let (from_bndrs, to_bndrs) = unzip bndrs
527
528 ; let from_bndr_tys = map idType from_bndrs -- Types ty
529
530
531 -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
532 ; expr' <- dsInnerMonadComp stmts from_bndrs return_op
533
534 -- Work out what arguments should be supplied to that expression: i.e. is an extraction
535 -- function required? If so, create that desugared function and add to arguments
536 ; usingExpr' <- dsLExpr using
537 ; usingArgs' <- case by of
538 Nothing -> return [expr']
539 Just by_e -> do { by_e' <- dsLExpr by_e
540 ; lam' <- matchTuple from_bndrs by_e'
541 ; return [lam', expr'] }
542
543 -- Generate the expressions to build the grouped list
544 -- Build a pattern that ensures the consumer binds into the NEW binders,
545 -- which hold monads rather than single values
546 ; let tup_n_ty' = mkBigCoreVarTupTy to_bndrs
547
548 ; body <- dsMcStmts stmts_rest
549 ; n_tup_var' <- newSysLocalDs Many n_tup_ty'
550 ; tup_n_var' <- newSysLocalDs Many tup_n_ty'
551 ; tup_n_expr' <- mkMcUnzipM form fmap_op n_tup_var' from_bndr_tys
552 ; us <- newUniqueSupply
553 ; let rhs' = mkApps usingExpr' usingArgs'
554 body' = mkTupleCase us to_bndrs body tup_n_var' tup_n_expr'
555
556 ; dsSyntaxExpr bind_op [rhs', Lam n_tup_var' body'] }
557
558 -- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel
559 -- statements, for example:
560 --
561 -- [ body | qs1 | qs2 | qs3 ]
562 -- -> [ body | (bndrs1, (bndrs2, bndrs3))
563 -- <- [bndrs1 | qs1] `mzip` ([bndrs2 | qs2] `mzip` [bndrs3 | qs3]) ]
564 --
565 -- where `mzip` has type
566 -- mzip :: forall a b. m a -> m b -> m (a,b)
567 -- NB: we need a polymorphic mzip because we call it several times
568
569 dsMcStmt (ParStmt bind_ty blocks mzip_op bind_op) stmts_rest
570 = do { exps_w_tys <- mapM ds_inner blocks -- Pairs (exp :: m ty, ty)
571 ; mzip_op' <- dsExpr mzip_op
572
573 ; let -- The pattern variables
574 pats = [ mkBigLHsVarPatTupId bs | ParStmtBlock _ _ bs _ <- blocks]
575 -- Pattern with tuples of variables
576 -- [v1,v2,v3] => (v1, (v2, v3))
577 pat = foldr1 (\p1 p2 -> mkLHsPatTup [p1, p2]) pats
578 (rhs, _) = foldr1 (\(e1,t1) (e2,t2) ->
579 (mkApps mzip_op' [Type t1, Type t2, e1, e2],
580 mkBoxedTupleTy [t1,t2]))
581 exps_w_tys
582
583 ; dsMcBindStmt pat rhs bind_op Nothing bind_ty stmts_rest }
584 where
585 ds_inner :: ParStmtBlock GhcTc GhcTc -> DsM (CoreExpr, Type)
586 ds_inner (ParStmtBlock _ stmts bndrs return_op)
587 = do { exp <- dsInnerMonadComp stmts bndrs return_op
588 ; return (exp, mkBigCoreVarTupTy bndrs) }
589
590 dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
591
592
593 matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
594 -- (matchTuple [a,b,c] body)
595 -- returns the Core term
596 -- \x. case x of (a,b,c) -> body
597 matchTuple ids body
598 = do { us <- newUniqueSupply
599 ; tup_id <- newSysLocalDs Many (mkBigCoreVarTupTy ids)
600 ; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
601
602 -- general `rhs' >>= \pat -> stmts` desugaring where `rhs'` is already a
603 -- desugared `CoreExpr`
604 dsMcBindStmt :: LPat GhcTc
605 -> CoreExpr -- ^ the desugared rhs of the bind statement
606 -> SyntaxExpr GhcTc
607 -> Maybe (SyntaxExpr GhcTc)
608 -> Type -- ^ S in (>>=) :: Q -> (R -> S) -> T
609 -> [ExprLStmt GhcTc]
610 -> DsM CoreExpr
611 dsMcBindStmt pat rhs' bind_op fail_op res1_ty stmts
612 = do { body <- dsMcStmts stmts
613 ; var <- selectSimpleMatchVarL Many pat
614 ; match <- matchSinglePatVar var Nothing (StmtCtxt (HsDoStmt (DoExpr Nothing))) pat
615 res1_ty (cantFailMatchResult body)
616 ; match_code <- dsHandleMonadicFailure MonadComp pat match fail_op
617 ; dsSyntaxExpr bind_op [rhs', Lam var match_code] }
618
619 -- Desugar nested monad comprehensions, for example in `then..` constructs
620 -- dsInnerMonadComp quals [a,b,c] ret_op
621 -- returns the desugaring of
622 -- [ (a,b,c) | quals ]
623
624 dsInnerMonadComp :: [ExprLStmt GhcTc]
625 -> [Id] -- Return a tuple of these variables
626 -> SyntaxExpr GhcTc -- The monomorphic "return" operator
627 -> DsM CoreExpr
628 dsInnerMonadComp stmts bndrs ret_op
629 = dsMcStmts (stmts ++
630 [noLocA (LastStmt noExtField (mkBigLHsVarTupId bndrs) Nothing ret_op)])
631
632
633 -- The `unzip` function for `GroupStmt` in a monad comprehensions
634 --
635 -- unzip :: m (a,b,..) -> (m a,m b,..)
636 -- unzip m_tuple = ( liftM selN1 m_tuple
637 -- , liftM selN2 m_tuple
638 -- , .. )
639 --
640 -- mkMcUnzipM fmap ys [t1, t2]
641 -- = ( fmap (selN1 :: (t1, t2) -> t1) ys
642 -- , fmap (selN2 :: (t1, t2) -> t2) ys )
643
644 mkMcUnzipM :: TransForm
645 -> HsExpr GhcTc -- fmap
646 -> Id -- Of type n (a,b,c)
647 -> [Type] -- [a,b,c] (not representation-polymorphic)
648 -> DsM CoreExpr -- Of type (n a, n b, n c)
649 mkMcUnzipM ThenForm _ ys _
650 = return (Var ys) -- No unzipping to do
651
652 mkMcUnzipM _ fmap_op ys elt_tys
653 = do { fmap_op' <- dsExpr fmap_op
654 ; xs <- mapM (newSysLocalDs Many) elt_tys
655 ; let tup_ty = mkBigCoreTupTy elt_tys
656 ; tup_xs <- newSysLocalDs Many tup_ty
657
658 ; let mk_elt i = mkApps fmap_op' -- fmap :: forall a b. (a -> b) -> n a -> n b
659 [ Type tup_ty, Type (getNth elt_tys i)
660 , mk_sel i, Var ys]
661
662 mk_sel n = Lam tup_xs $
663 mkTupleSelector xs (getNth xs n) tup_xs (Var tup_xs)
664
665 ; return (mkBigCoreTup (map mk_elt [0..length elt_tys - 1])) }