1ecab67e10fbc580e8e79d854152c004db1b2b8d
[ghc-hetmet.git] / compiler / deSugar / DsListComp.lhs
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 \begin{code}
9 {-# LANGUAGE NamedFieldPuns #-}
10 {-# OPTIONS -fno-warn-incomplete-patterns #-}
11 -- The above warning supression flag is a temporary kludge.
12 -- While working on this module you are encouraged to remove it and fix
13 -- any warnings in the module. See
14 --     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
15 -- for details
16
17 module DsListComp ( dsListComp, dsPArrComp, dsMonadComp ) where
18
19 #include "HsVersions.h"
20
21 import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLocalBinds )
22
23 import HsSyn
24 import TcHsSyn
25 import CoreSyn
26 import MkCore
27
28 import DsMonad          -- the monadery used in the desugarer
29 import DsUtils
30
31 import DynFlags
32 import CoreUtils
33 import Id
34 import Type
35 import TysWiredIn
36 import Match
37 import PrelNames
38 import SrcLoc
39 import Outputable
40 import FastString
41 import TcType
42 \end{code}
43
44 List comprehensions may be desugared in one of two ways: ``ordinary''
45 (as you would expect if you read SLPJ's book) and ``with foldr/build
46 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
47
48 There will be at least one ``qualifier'' in the input.
49
50 \begin{code}
51 dsListComp :: [LStmt Id] 
52            -> Type              -- Type of entire list 
53            -> DsM CoreExpr
54 dsListComp lquals res_ty = do 
55     dflags <- getDOptsDs
56     let quals = map unLoc lquals
57         [elt_ty] = tcTyConAppArgs res_ty
58     
59     if not (dopt Opt_EnableRewriteRules dflags) || dopt 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 :: ([LStmt Id], [Id]) -> DsM (CoreExpr, Type)
85 dsInnerListComp (stmts, bndrs) = do
86   = do { expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTup bndrs)]) 
87                            bndrs_tuple_type
88        ; return (expr, bndrs_tuple_type) }
89   where
90     bndrs_tuple_type = mkBigCoreVarTupTy bndrs
91         
92 -- This function factors out commonality between the desugaring strategies for TransformStmt.
93 -- Given such a statement it gives you back an expression representing how to compute the transformed
94 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
95 dsTransformStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
96 dsTransformStmt (TransformStmt stmts binders usingExpr maybeByExpr _ _)
97  = do { (expr, binders_tuple_type) <- dsInnerListComp (stmts, binders)
98       ; usingExpr' <- dsLExpr usingExpr
99     
100       ; using_args <-
101           case maybeByExpr of
102             Nothing -> return [expr]
103             Just byExpr -> do
104                 byExpr' <- dsLExpr byExpr
105                 
106                 us <- newUniqueSupply
107                 [tuple_binder] <- newSysLocalsDs [binders_tuple_type]
108                 let byExprWrapper = mkTupleCase us binders byExpr' tuple_binder (Var tuple_binder)
109                 
110                 return [Lam tuple_binder byExprWrapper, expr]
111
112       ; let inner_list_expr = mkApps usingExpr' ((Type binders_tuple_type) : using_args)
113             pat = mkBigLHsVarPatTup binders
114       ; return (inner_list_expr, pat) }
115     
116 -- This function factors out commonality between the desugaring strategies for GroupStmt.
117 -- Given such a statement it gives you back an expression representing how to compute the transformed
118 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
119 dsGroupStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
120 dsGroupStmt (GroupStmt stmts binderMap by using _ _ _) = do
121     let (fromBinders, toBinders) = unzip binderMap
122         
123         fromBindersTypes = map idType fromBinders
124         toBindersTypes = map idType toBinders
125         
126         toBindersTupleType = mkBigCoreTupTy toBindersTypes
127     
128     -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
129     (expr, from_tup_ty) <- dsInnerListComp (stmts, fromBinders)
130     
131     -- Work out what arguments should be supplied to that expression: i.e. is an extraction
132     -- function required? If so, create that desugared function and add to arguments
133     usingExpr' <- dsLExpr (either id noLoc using)
134     usingArgs <- case by of
135                    Nothing   -> return [expr]
136                    Just by_e -> do { by_e' <- dsLExpr by_e
137                                    ; us <- newUniqueSupply
138                                    ; [from_tup_id] <- newSysLocalsDs [from_tup_ty]
139                                    ; let by_wrap = mkTupleCase us fromBinders by_e' 
140                                                    from_tup_id (Var from_tup_id)
141                                    ; return [Lam from_tup_id by_wrap, expr] }
142     
143     -- Create an unzip function for the appropriate arity and element types and find "map"
144     (unzip_fn, unzip_rhs) <- mkUnzipBind fromBindersTypes
145     map_id <- dsLookupGlobalId mapName
146
147     -- Generate the expressions to build the grouped list
148     let -- First we apply the grouping function to the inner list
149         inner_list_expr = mkApps usingExpr' ((Type from_tup_ty) : usingArgs)
150         -- Then we map our "unzip" across it to turn the lists of tuples into tuples of lists
151         -- We make sure we instantiate the type variable "a" to be a list of "from" tuples and
152         -- the "b" to be a tuple of "to" lists!
153         unzipped_inner_list_expr = mkApps (Var map_id) 
154             [Type (mkListTy from_tup_ty), Type toBindersTupleType, Var unzip_fn, inner_list_expr]
155         -- Then finally we bind the unzip function around that expression
156         bound_unzipped_inner_list_expr = Let (Rec [(unzip_fn, unzip_rhs)]) unzipped_inner_list_expr
157     
158     -- Build a pattern that ensures the consumer binds into the NEW binders, which hold lists rather than single values
159     let pat = mkBigLHsVarPatTup toBinders
160     return (bound_unzipped_inner_list_expr, pat)
161     
162 \end{code}
163
164 %************************************************************************
165 %*                                                                      *
166 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
167 %*                                                                      *
168 %************************************************************************
169
170 Just as in Phil's chapter~7 in SLPJ, using the rules for
171 optimally-compiled list comprehensions.  This is what Kevin followed
172 as well, and I quite happily do the same.  The TQ translation scheme
173 transforms a list of qualifiers (either boolean expressions or
174 generators) into a single expression which implements the list
175 comprehension.  Because we are generating 2nd-order polymorphic
176 lambda-calculus, calls to NIL and CONS must be applied to a type
177 argument, as well as their usual value arguments.
178 \begin{verbatim}
179 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
180
181 (Rule C)
182 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
183
184 (Rule B)
185 TQ << [ e | b , qs ] ++ L >> =
186     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
187
188 (Rule A')
189 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
190   letrec
191     h = \ u1 ->
192           case u1 of
193             []        ->  TE << L2 >>
194             (u2 : u3) ->
195                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
196                     [] (h u3)
197   in
198     h ( TE << L1 >> )
199
200 "h", "u1", "u2", and "u3" are new variables.
201 \end{verbatim}
202
203 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
204 is the TE translation scheme.  Note that we carry around the @L@ list
205 already desugared.  @dsListComp@ does the top TE rule mentioned above.
206
207 To the above, we add an additional rule to deal with parallel list
208 comprehensions.  The translation goes roughly as follows:
209      [ e | p1 <- e11, let v1 = e12, p2 <- e13
210          | q1 <- e21, let v2 = e22, q2 <- e23]
211      =>
212      [ e | ((x1, .., xn), (y1, ..., ym)) <-
213                zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
214                    [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
215 where (x1, .., xn) are the variables bound in p1, v1, p2
216       (y1, .., ym) are the variables bound in q1, v2, q2
217
218 In the translation below, the ParStmt branch translates each parallel branch
219 into a sub-comprehension, and desugars each independently.  The resulting lists
220 are fed to a zip function, we create a binding for all the variables bound in all
221 the comprehensions, and then we hand things off the the desugarer for bindings.
222 The zip function is generated here a) because it's small, and b) because then we
223 don't have to deal with arbitrary limits on the number of zip functions in the
224 prelude, nor which library the zip function came from.
225 The introduced tuples are Boxed, but only because I couldn't get it to work
226 with the Unboxed variety.
227
228 \begin{code}
229
230 deListComp :: [Stmt Id] -> CoreExpr -> DsM CoreExpr
231
232 deListComp [] _ = panic "deListComp"
233
234 deListComp (LastStmt body _ : quals) list 
235   =     -- Figure 7.4, SLPJ, p 135, rule C above
236     ASSERT( null quals )
237     do { core_body <- dsLExpr body
238        ; return (mkConsExpr (exprType core_body) core_body list) }
239
240         -- Non-last: must be a guard
241 deListComp (ExprStmt guard _ _ _ : quals) list = do  -- rule B above
242     core_guard <- dsLExpr guard
243     core_rest <- deListComp quals list
244     return (mkIfThenElse core_guard core_rest list)
245
246 -- [e | let B, qs] = let B in [e | qs]
247 deListComp (LetStmt binds : quals) list = do
248     core_rest <- deListComp quals list
249     dsLocalBinds binds core_rest
250
251 deListComp (stmt@(TransformStmt {}) : quals) list = do
252     (inner_list_expr, pat) <- dsTransformStmt stmt
253     deBindComp pat inner_list_expr quals list
254
255 deListComp (stmt@(GroupStmt {}) : quals) list = do
256     (inner_list_expr, pat) <- dsGroupStmt stmt
257     deBindComp pat inner_list_expr quals list
258
259 deListComp (BindStmt pat list1 _ _ : quals) core_list2 = do -- rule A' above
260     core_list1 <- dsLExpr list1
261     deBindComp pat core_list1 quals core_list2
262
263 deListComp (ParStmt stmtss_w_bndrs _ _ _ : quals) list
264   = do
265     exps_and_qual_tys <- mapM dsInnerListComp stmtss_w_bndrs
266     let (exps, qual_tys) = unzip exps_and_qual_tys
267     
268     (zip_fn, zip_rhs) <- mkZipBind qual_tys
269
270         -- Deal with [e | pat <- zip l1 .. ln] in example above
271     deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps)) 
272                    quals list
273
274   where 
275         bndrs_s = map snd stmtss_w_bndrs
276
277         -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
278         pat  = mkBigLHsPatTup pats
279         pats = map mkBigLHsVarPatTup bndrs_s
280 \end{code}
281
282
283 \begin{code}
284 deBindComp :: OutPat Id
285            -> CoreExpr
286            -> [Stmt Id]
287            -> CoreExpr
288            -> DsM (Expr Id)
289 deBindComp pat core_list1 quals core_list2 = do
290     let
291         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
292
293         -- u1_ty is a [alpha] type, and u2_ty = alpha
294         u2_ty = hsLPatType pat
295
296         res_ty = exprType core_list2
297         h_ty   = u1_ty `mkFunTy` res_ty
298         
299     [h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]
300
301     -- the "fail" value ...
302     let
303         core_fail   = App (Var h) (Var u3)
304         letrec_body = App (Var h) core_list1
305         
306     rest_expr <- deListComp quals core_fail
307     core_match <- matchSimply (Var u2) (StmtCtxt ListComp) pat rest_expr core_fail      
308     
309     let
310         rhs = Lam u1 $
311               Case (Var u1) u1 res_ty
312                    [(DataAlt nilDataCon,  [],       core_list2),
313                     (DataAlt consDataCon, [u2, u3], core_match)]
314                         -- Increasing order of tag
315             
316     return (Let (Rec [(h, rhs)]) letrec_body)
317 \end{code}
318
319 %************************************************************************
320 %*                                                                      *
321 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
322 %*                                                                      *
323 %************************************************************************
324
325 @dfListComp@ are the rules used with foldr/build turned on:
326
327 \begin{verbatim}
328 TE[ e | ]            c n = c e n
329 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
330 TE[ e | p <- l , q ] c n = let 
331                                 f = \ x b -> case x of
332                                                   p -> TE[ e | q ] c b
333                                                   _ -> b
334                            in
335                            foldr f n l
336 \end{verbatim}
337
338 \begin{code}
339 dfListComp :: Id -> Id -- 'c' and 'n'
340         -> [Stmt Id]   -- the rest of the qual's
341         -> DsM CoreExpr
342
343 dfListComp _ _ [] = panic "dfListComp"
344
345 dfListComp c_id n_id (LastStmt body _ : quals) 
346   = ASSERT( null quals )
347     do { core_body <- dsLExpr body
348        ; return (mkApps (Var c_id) [core_body, Var n_id]) }
349
350         -- Non-last: must be a guard
351 dfListComp c_id n_id (ExprStmt guard _ _ _  : quals) = do
352     core_guard <- dsLExpr guard
353     core_rest <- dfListComp c_id n_id quals
354     return (mkIfThenElse core_guard core_rest (Var n_id))
355
356 dfListComp c_id n_id (LetStmt binds : quals) = do
357     -- new in 1.3, local bindings
358     core_rest <- dfListComp c_id n_id quals
359     dsLocalBinds binds core_rest
360
361 dfListComp c_id n_id (stmt@(TransformStmt {}) : quals) = do
362     (inner_list_expr, pat) <- dsTransformStmt stmt
363     -- Anyway, we bind the newly transformed list via the generic binding function
364     dfBindComp c_id n_id (pat, inner_list_expr) quals 
365
366 dfListComp c_id n_id (stmt@(GroupStmt {}) : quals) = do
367     (inner_list_expr, pat) <- dsGroupStmt stmt
368     -- Anyway, we bind the newly grouped list via the generic binding function
369     dfBindComp c_id n_id (pat, inner_list_expr) quals 
370     
371 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) = do
372     -- evaluate the two lists
373     core_list1 <- dsLExpr list1
374     
375     -- Do the rest of the work in the generic binding builder
376     dfBindComp c_id n_id (pat, core_list1) quals
377                
378 dfBindComp :: Id -> Id          -- 'c' and 'n'
379        -> (LPat Id, CoreExpr)
380            -> [Stmt Id]                 -- the rest of the qual's
381            -> DsM CoreExpr
382 dfBindComp c_id n_id (pat, core_list1) quals = do
383     -- find the required type
384     let x_ty   = hsLPatType pat
385         b_ty   = idType n_id
386
387     -- create some new local id's
388     [b, x] <- newSysLocalsDs [b_ty, x_ty]
389
390     -- build rest of the comprehesion
391     core_rest <- dfListComp c_id b quals
392
393     -- build the pattern match
394     core_expr <- matchSimply (Var x) (StmtCtxt ListComp)
395                 pat core_rest (Var b)
396
397     -- now build the outermost foldr, and return
398     mkFoldrExpr x_ty b_ty (mkLams [x, b] core_expr) (Var n_id) core_list1
399 \end{code}
400
401 %************************************************************************
402 %*                                                                      *
403 \subsection[DsFunGeneration]{Generation of zip/unzip functions for use in desugaring}
404 %*                                                                      *
405 %************************************************************************
406
407 \begin{code}
408
409 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
410 -- mkZipBind [t1, t2] 
411 -- = (zip, \as1:[t1] as2:[t2] 
412 --         -> case as1 of 
413 --              [] -> []
414 --              (a1:as'1) -> case as2 of
415 --                              [] -> []
416 --                              (a2:as'2) -> (a1, a2) : zip as'1 as'2)]
417
418 mkZipBind elt_tys = do
419     ass  <- mapM newSysLocalDs  elt_list_tys
420     as'  <- mapM newSysLocalDs  elt_tys
421     as's <- mapM newSysLocalDs  elt_list_tys
422     
423     zip_fn <- newSysLocalDs zip_fn_ty
424     
425     let inner_rhs = mkConsExpr elt_tuple_ty 
426                         (mkBigCoreVarTup as')
427                         (mkVarApps (Var zip_fn) as's)
428         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
429     
430     return (zip_fn, mkLams ass zip_body)
431   where
432     elt_list_tys      = map mkListTy elt_tys
433     elt_tuple_ty      = mkBigCoreTupTy elt_tys
434     elt_tuple_list_ty = mkListTy elt_tuple_ty
435     
436     zip_fn_ty         = mkFunTys elt_list_tys elt_tuple_list_ty
437
438     mk_case (as, a', as') rest
439           = Case (Var as) as elt_tuple_list_ty
440                   [(DataAlt nilDataCon,  [],        mkNilExpr elt_tuple_ty),
441                    (DataAlt consDataCon, [a', as'], rest)]
442                         -- Increasing order of tag
443             
444             
445 mkUnzipBind :: [Type] -> DsM (Id, CoreExpr)
446 -- mkUnzipBind [t1, t2] 
447 -- = (unzip, \ys :: [(t1, t2)] -> foldr (\ax :: (t1, t2) axs :: ([t1], [t2])
448 --     -> case ax of
449 --      (x1, x2) -> case axs of
450 --                (xs1, xs2) -> (x1 : xs1, x2 : xs2))
451 --      ([], [])
452 --      ys)
453 -- 
454 -- We use foldr here in all cases, even if rules are turned off, because we may as well!
455 mkUnzipBind elt_tys = do
456     ax  <- newSysLocalDs elt_tuple_ty
457     axs <- newSysLocalDs elt_list_tuple_ty
458     ys  <- newSysLocalDs elt_tuple_list_ty
459     xs  <- mapM newSysLocalDs elt_tys
460     xss <- mapM newSysLocalDs elt_list_tys
461     
462     unzip_fn <- newSysLocalDs unzip_fn_ty
463
464     [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
465
466     let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
467         
468         concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
469         tupled_concat_expression = mkBigCoreTup concat_expressions
470         
471         folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
472         folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
473         folder_body = mkLams [ax, axs] folder_body_outer_case
474         
475     unzip_body <- mkFoldrExpr elt_tuple_ty elt_list_tuple_ty folder_body nil_tuple (Var ys)
476     return (unzip_fn, mkLams [ys] unzip_body)
477   where
478     elt_tuple_ty       = mkBigCoreTupTy elt_tys
479     elt_tuple_list_ty  = mkListTy elt_tuple_ty
480     elt_list_tys       = map mkListTy elt_tys
481     elt_list_tuple_ty  = mkBigCoreTupTy elt_list_tys
482     
483     unzip_fn_ty        = elt_tuple_list_ty `mkFunTy` elt_list_tuple_ty
484             
485     mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
486 \end{code}
487
488 %************************************************************************
489 %*                                                                      *
490 \subsection[DsPArrComp]{Desugaring of array comprehensions}
491 %*                                                                      *
492 %************************************************************************
493
494 \begin{code}
495
496 -- entry point for desugaring a parallel array comprehension
497 --
498 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
499 --
500 dsPArrComp :: [Stmt Id] 
501             -> DsM CoreExpr
502
503 -- Special case for parallel comprehension
504 dsPArrComp (ParStmt qss _ _ _ : quals) = dePArrParComp qss quals
505
506 -- Special case for simple generators:
507 --
508 --  <<[:e' | p <- e, qs:]>> = <<[: e' | qs :]>> p e
509 --
510 -- if matching again p cannot fail, or else
511 --
512 --  <<[:e' | p <- e, qs:]>> = 
513 --    <<[:e' | qs:]>> p (filterP (\x -> case x of {p -> True; _ -> False}) e)
514 --
515 dsPArrComp (BindStmt p e _ _ : qs) = do
516     filterP <- dsLookupDPHId filterPName
517     ce <- dsLExpr e
518     let ety'ce  = parrElemType ce
519         false   = Var falseDataConId
520         true    = Var trueDataConId
521     v <- newSysLocalDs ety'ce
522     pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
523     let gen | isIrrefutableHsPat p = ce
524             | otherwise            = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
525     dePArrComp qs p gen
526
527 dsPArrComp qs = do -- no ParStmt in `qs'
528     sglP <- dsLookupDPHId singletonPName
529     let unitArray = mkApps (Var sglP) [Type unitTy, mkCoreTup []]
530     dePArrComp qs (noLoc $ WildPat unitTy) unitArray
531
532
533
534 -- the work horse
535 --
536 dePArrComp :: [Stmt Id] 
537            -> LPat Id           -- the current generator pattern
538            -> CoreExpr          -- the current generator expression
539            -> DsM CoreExpr
540
541 dePArrComp [] _ _ = panic "dePArrComp"
542
543 --
544 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
545 --
546 dePArrComp (LastStmt e' _ : quals) pa cea
547   = ASSERT( null quals )
548     do { mapP <- dsLookupDPHId mapPName
549        ; let ty = parrElemType cea
550        ; (clam, ty'e') <- deLambda ty pa e'
551        ; return $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea] }
552 --
553 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
554 --
555 dePArrComp (ExprStmt b _ _ _ : qs) pa cea = do
556     filterP <- dsLookupDPHId filterPName
557     let ty = parrElemType cea
558     (clam,_) <- deLambda ty pa b
559     dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
560
561 --
562 --  <<[:e' | p <- e, qs:]>> pa ea =
563 --    let ef = \pa -> e
564 --    in
565 --    <<[:e' | qs:]>> (pa, p) (crossMap ea ef)
566 --
567 -- if matching again p cannot fail, or else
568 --
569 --  <<[:e' | p <- e, qs:]>> pa ea = 
570 --    let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
571 --    in
572 --    <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
573 --
574 dePArrComp (BindStmt p e _ _ : qs) pa cea = do
575     filterP <- dsLookupDPHId filterPName
576     crossMapP <- dsLookupDPHId crossMapPName
577     ce <- dsLExpr e
578     let ety'cea = parrElemType cea
579         ety'ce  = parrElemType ce
580         false   = Var falseDataConId
581         true    = Var trueDataConId
582     v <- newSysLocalDs ety'ce
583     pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
584     let cef | isIrrefutableHsPat p = ce
585             | otherwise            = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
586     (clam, _) <- mkLambda ety'cea pa cef
587     let ety'cef = ety'ce                    -- filter doesn't change the element type
588         pa'     = mkLHsPatTup [pa, p]
589
590     dePArrComp qs pa' (mkApps (Var crossMapP) 
591                                  [Type ety'cea, Type ety'cef, cea, clam])
592 --
593 --  <<[:e' | let ds, qs:]>> pa ea = 
594 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n)) 
595 --                    (mapP (\v@pa -> let ds in (v, (x_1, ..., x_n))) ea)
596 --  where
597 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
598 --
599 dePArrComp (LetStmt ds : qs) pa cea = do
600     mapP <- dsLookupDPHId mapPName
601     let xs     = collectLocalBinders ds
602         ty'cea = parrElemType cea
603     v <- newSysLocalDs ty'cea
604     clet <- dsLocalBinds ds (mkCoreTup (map Var xs))
605     let'v <- newSysLocalDs (exprType clet)
606     let projBody = mkCoreLet (NonRec let'v clet) $ 
607                    mkCoreTup [Var v, Var let'v]
608         errTy    = exprType projBody
609         errMsg   = ptext (sLit "DsListComp.dePArrComp: internal error!")
610     cerr <- mkErrorAppDs pAT_ERROR_ID errTy errMsg
611     ccase <- matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr
612     let pa'    = mkLHsPatTup [pa, mkLHsPatTup (map nlVarPat xs)]
613         proj   = mkLams [v] ccase
614     dePArrComp qs pa' (mkApps (Var mapP) 
615                                    [Type ty'cea, Type errTy, proj, cea])
616 --
617 -- The parser guarantees that parallel comprehensions can only appear as
618 -- singeltons qualifier lists, which we already special case in the caller.
619 -- So, encountering one here is a bug.
620 --
621 dePArrComp (ParStmt _ _ _ _ : _) _ _ = 
622   panic "DsListComp.dePArrComp: malformed comprehension AST"
623
624 --  <<[:e' | qs | qss:]>> pa ea = 
625 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n)) 
626 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
627 --    where
628 --      {x_1, ..., x_n} = DV (qs)
629 --
630 dePArrParComp :: [([LStmt Id], [Id])] -> [Stmt Id] -> DsM CoreExpr
631 dePArrParComp qss quals = do
632     (pQss, ceQss) <- deParStmt qss
633     dePArrComp quals pQss ceQss
634   where
635     deParStmt []             =
636       -- empty parallel statement lists have no source representation
637       panic "DsListComp.dePArrComp: Empty parallel list comprehension"
638     deParStmt ((qs, xs):qss) = do        -- first statement
639       let res_expr = mkLHsVarTuple xs
640       cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
641       parStmts qss (mkLHsVarPatTup xs) cqs
642     ---
643     parStmts []             pa cea = return (pa, cea)
644     parStmts ((qs, xs):qss) pa cea = do  -- subsequent statements (zip'ed)
645       zipP <- dsLookupDPHId zipPName
646       let pa'      = mkLHsPatTup [pa, mkLHsVarPatTup xs]
647           ty'cea   = parrElemType cea
648           res_expr = mkLHsVarTuple xs
649       cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
650       let ty'cqs = parrElemType cqs
651           cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
652       parStmts qss pa' cea'
653
654 -- generate Core corresponding to `\p -> e'
655 --
656 deLambda :: Type                        -- type of the argument
657           -> LPat Id                    -- argument pattern
658           -> LHsExpr Id                 -- body
659           -> DsM (CoreExpr, Type)
660 deLambda ty p e =
661     mkLambda ty p =<< dsLExpr e
662
663 -- generate Core for a lambda pattern match, where the body is already in Core
664 --
665 mkLambda :: Type                        -- type of the argument
666          -> LPat Id                     -- argument pattern
667          -> CoreExpr                    -- desugared body
668          -> DsM (CoreExpr, Type)
669 mkLambda ty p ce = do
670     v <- newSysLocalDs ty
671     let errMsg = ptext (sLit "DsListComp.deLambda: internal error!")
672         ce'ty  = exprType ce
673     cerr <- mkErrorAppDs pAT_ERROR_ID ce'ty errMsg
674     res <- matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr
675     return (mkLams [v] res, ce'ty)
676
677 -- obtain the element type of the parallel array produced by the given Core
678 -- expression
679 --
680 parrElemType   :: CoreExpr -> Type
681 parrElemType e  = 
682   case splitTyConApp_maybe (exprType e) of
683     Just (tycon, [ty]) | tycon == parrTyCon -> ty
684     _                                                     -> panic
685       "DsListComp.parrElemType: not a parallel array type"
686 \end{code}
687
688 Translation for monad comprehensions
689
690 \begin{code}
691
692 -- | Keep the "context" of a monad comprehension in a small data type to avoid
693 -- some boilerplate...
694 data DsMonadComp = DsMonadComp
695     { mc_return :: Either (SyntaxExpr Id) (Expr CoreBndr)
696     , mc_body   :: LHsExpr Id
697     , mc_m_ty   :: Type
698     }
699
700 --
701 -- Entry point for monad comprehension desugaring
702 --
703 dsMonadComp :: [LStmt Id]       -- the statements
704             -> Type             -- the final type
705             -> DsM CoreExpr
706 dsMonadComp stmts res_ty
707   = dsMcStmts stmts (DsMonadComp (Left return_op) body m_ty)
708   where
709     (m_ty, _) = tcSplitAppTy res_ty
710
711
712 dsMcStmts :: [LStmt Id]
713           -> DsMonadComp
714           -> DsM CoreExpr
715
716 -- No statements left for desugaring. Desugar the body after calling "return"
717 -- on it.
718 dsMcStmts [] DsMonadComp { mc_return, mc_body }
719   = case mc_return of
720          Left ret   -> dsLExpr $ noLoc ret `nlHsApp` mc_body
721          Right ret' -> do
722              { body' <- dsLExpr mc_body
723              ; return $ mkApps ret' [body'] }
724
725 -- Otherwise desugar each statement step by step
726 dsMcStmts ((L loc stmt) : lstmts) mc
727   = putSrcSpanDs loc (dsMcStmt stmt lstmts mc)
728
729
730 dsMcStmt :: Stmt Id -> [LStmt Id] -> DsM CoreExpr
731
732 dsMcStmt (LastStmt body ret_op) stmts
733   = ASSERT( null stmts )
734     do { body' <- dsLExpr body
735        ; ret_op' <- dsExpr ret_op
736        ; return (App ret_op' body') }
737
738 --   [ .. | let binds, stmts ]
739 dsMcStmt (LetStmt binds) stmts 
740   = do { rest <- dsMcStmts stmts
741        ; dsLocalBinds binds rest }
742
743 --   [ .. | a <- m, stmts ]
744 dsMcStmt (BindStmt pat rhs bind_op fail_op) stmts
745   = do { rhs' <- dsLExpr rhs
746        ; dsMcBindStmt pat rhs' bind_op fail_op stmts }
747
748 -- Apply `guard` to the `exp` expression
749 --
750 --   [ .. | exp, stmts ]
751 --
752 dsMcStmt (ExprStmt exp then_exp guard_exp _) stmts 
753   = do { exp'       <- dsLExpr exp
754        ; guard_exp' <- dsExpr guard_exp
755        ; then_exp'  <- dsExpr then_exp
756        ; rest       <- dsMcStmts stmts
757        ; return $ mkApps then_exp' [ mkApps guard_exp' [exp']
758                                    , rest ] }
759
760 -- Transform statements desugar like this:
761 --
762 --   [ .. | qs, then f by e ]  ->  f (\q_v -> e) [| qs |]
763 --
764 -- where [| qs |] is the desugared inner monad comprehenion generated by the
765 -- statements `qs`.
766 dsMcStmt (TransformStmt stmts binders usingExpr maybeByExpr return_op bind_op) stmts_rest
767   = do { expr <- dsInnerMonadComp stmts binders return_op
768        ; let binders_tup_type = mkBigCoreTupTy $ map idType binders
769        ; usingExpr' <- dsLExpr usingExpr
770        ; using_args <- case maybeByExpr of
771             Nothing -> return [expr]
772             Just byExpr -> do
773                 byExpr' <- dsLExpr byExpr
774                 us <- newUniqueSupply
775                 tup_binder <- newSysLocalDs binders_tup_type
776                 let byExprWrapper = mkTupleCase us binders byExpr' tup_binder (Var tup_binder)
777                 return [Lam tup_binder byExprWrapper, expr]
778
779        ; let pat = mkBigLHsVarPatTup binders
780              rhs = mkApps usingExpr' ((Type binders_tup_type) : using_args)
781
782        ; dsMcBindStmt pat rhs bind_op noSyntaxExpr stmts_rest }
783
784 -- Group statements desugar like this:
785 --
786 --   [| (q, then group by e using f); rest |]
787 --   --->  f {qt} (\qv -> e) [| q; return qv |] >>= \ n_tup -> 
788 --         case unzip n_tup of qv -> [| rest |]
789 --
790 -- where   variables (v1:t1, ..., vk:tk) are bound by q
791 --         qv = (v1, ..., vk)
792 --         qt = (t1, ..., tk)
793 --         (>>=) :: m2 a -> (a -> m3 b) -> m3 b
794 --         f :: forall a. (a -> t) -> m1 a -> m2 (n a)
795 --         n_tup :: n qt
796 --         unzip :: n qt -> (n t1, ..., n tk)    (needs Functor n)
797 --
798 --   [| q, then group by e using f |]  ->  (f (\q_v -> e) [| q |]) >>= (return . (unzip q_v))
799 --
800 -- which is equal to
801 --
802 --   [| q, then group by e using f |]  ->  liftM (unzip q_v) (f (\q_v -> e) [| q |])
803 --
804 -- where unzip is of the form
805 --
806 --   unzip :: n (a,b,c,..) -> (n a,n b,n c,..)
807 --   unzip m_tuple = ( fmap selN1 m_tuple
808 --                   , fmap selN2 m_tuple
809 --                   , .. )
810 --     where selN1 (a,b,c,..) = a
811 --           selN2 (a,b,c,..) = b
812 --           ..
813 --
814 dsMcStmt (GroupStmt stmts binderMap by using return_op bind_op fmap_op) stmts_rest
815   = do { let (fromBinders, toBinders) = unzip binderMap
816              fromBindersTypes         = map idType fromBinders          -- Types ty
817              fromBindersTupleTy       = mkBigCoreTupTy fromBindersTypes
818              toBindersTypes           = map idType toBinders            -- Types (n ty)
819              toBindersTupleTy         = mkBigCoreTupTy toBindersTypes
820
821        -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
822        ; expr <- dsInnerMonadComp stmts fromBinders return_op
823
824        -- Work out what arguments should be supplied to that expression: i.e. is an extraction
825        -- function required? If so, create that desugared function and add to arguments
826        ; usingExpr' <- dsLExpr (either id noLoc using)
827        ; usingArgs <- case by of
828                         Nothing   -> return [expr]
829                         Just by_e -> do { by_e' <- dsLExpr by_e
830                                         ; lam <- matchTuple fromBinders by_e'
831                                         ; return [lam, expr] }
832
833        -- Create an unzip function for the appropriate arity and element types
834        ; fmap_op' <- dsExpr fmap_op
835        ; (unzip_fn, unzip_rhs) <- mkMcUnzipM fmap_op' m_ty fromBindersTypes
836
837        -- Generate the expressions to build the grouped list
838        -- Build a pattern that ensures the consumer binds into the NEW binders, 
839        -- which hold monads rather than single values
840        ; bind_op' <- dsExpr bind_op
841        ; let bind_ty = exprType bind_op'    -- m2 (n (a,b,c)) -> (n (a,b,c) -> r1) -> r2
842              n_tup_ty = funArgTy $ funArgTy $ funResultTy bind_ty
843
844        ; body      <- dsMcStmts stmts_rest
845        ; n_tup_var <- newSysLocalDs n_tup_ty
846        ; tup_n_var <- newSysLocalDs (mkBigCoreVarTupTy toBinders)
847        ; us        <- newUniqueSupply
848        ; let unzip_n_tup = Let (Rec [(unzip_fn, unzip_rhs)]) $
849                            App (Var unzip_fn) (Var n_tup_var)
850              -- unzip_n_tup :: (n a, n b, n c)
851              body' = mkTupleCase us toBinders body unzip_n_tup (Var tup_n_var)
852                    
853        ; return (mkApps bind_op' [rhs', Lam n_tup_var body']) }
854
855 -- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel
856 -- statements, for example:
857 --
858 --   [ body | qs1 | qs2 | qs3 ]
859 --     ->  [ body | (bndrs1, (bndrs2, bndrs3)) 
860 --                     <- [bndrs1 | qs1] `mzip` ([bndrs2 | qs2] `mzip` [bndrs3 | qs3]) ]
861 --
862 -- where `mzip` has type
863 --   mzip :: forall a b. m a -> m b -> m (a,b)
864 -- NB: we need a polymorphic mzip because we call it several times
865
866 dsMcStmt (ParStmt pairs mzip_op bind_op return_op) stmts_rest
867  = do  { exps <- mapM ds_inner pairs
868        ; let qual_tys = map (mkBigCoreVarTupTy . snd) pairs
869        ; mzip_op' <- dsExpr mzip_op
870        ; (zip_fn, zip_rhs) <- mkMcZipM mzip_op' (mc_m_ty mc) qual_tys
871
872        ; let -- The pattern variables
873              vars = map (mkBigLHsVarPatTup . snd) pairs
874              -- Pattern with tuples of variables
875              -- [v1,v2,v3]  =>  (v1, (v2, v3))
876              pat = foldr (\tn tm -> mkBigLHsPatTup [tn, tm]) (last vars) (init vars)
877              rhs = Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps)
878
879        ; dsMcBindStmt pat rhs bind_op noSyntaxExpr stmts_rest }
880   where
881     ds_inner (stmts, bndrs) = dsInnerMonadComp stmts bndrs mono_ret_op
882        where 
883          mono_ret_op = HsWrap (WpTyApp (mkBigCoreVarTupTy bndrs)) return_op
884
885 dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
886
887
888 matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
889 -- (matchTuple [a,b,c] body)
890 --       returns the Core term
891 --  \x. case x of (a,b,c) -> body 
892 matchTuple ids body
893   = do { us <- newUniqueSupply
894        ; tup_id <- newSysLocalDs (mkBigLHsVarPatTup ids)
895        ; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
896
897
898 -- general `rhs' >>= \pat -> stmts` desugaring where `rhs'` is already a
899 -- desugared `CoreExpr`
900 dsMcBindStmt :: LPat Id
901              -> CoreExpr        -- ^ the desugared rhs of the bind statement
902              -> SyntaxExpr Id
903              -> SyntaxExpr Id
904              -> [LStmt Id]
905              -> DsM CoreExpr
906 dsMcBindStmt pat rhs' bind_op fail_op stmts
907   = do  { body     <- dsMcStmts stmts 
908         ; bind_op' <- dsExpr bind_op
909         ; var      <- selectSimpleMatchVarL pat
910         ; let bind_ty = exprType bind_op'       -- rhs -> (pat -> res1) -> res2
911               res1_ty = funResultTy (funArgTy (funResultTy bind_ty))
912         ; match <- matchSinglePat (Var var) (StmtCtxt DoExpr) pat
913                                   res1_ty (cantFailMatchResult body)
914         ; match_code <- handle_failure pat match fail_op
915         ; return (mkApps bind_op' [rhs', Lam var match_code]) }
916
917   where
918     -- In a monad comprehension expression, pattern-match failure just calls
919     -- the monadic `fail` rather than throwing an exception
920     handle_failure pat match fail_op
921       | matchCanFail match
922         = do { fail_op' <- dsExpr fail_op
923              ; fail_msg <- mkStringExpr (mk_fail_msg pat)
924              ; extractMatchResult match (App fail_op' fail_msg) }
925       | otherwise
926         = extractMatchResult match (error "It can't fail") 
927
928     mk_fail_msg :: Located e -> String
929     mk_fail_msg pat = "Pattern match failure in monad comprehension at " ++ 
930                       showSDoc (ppr (getLoc pat))
931
932 -- Desugar nested monad comprehensions, for example in `then..` constructs
933 --    dsInnerMonadComp quals [a,b,c] ret_op
934 -- returns the desugaring of 
935 --       [ (a,b,c) | quals ]
936
937 dsInnerMonadComp :: [LStmt Id]
938                  -> [Id]        -- Return a tuple of these variables
939                  -> LHsExpr Id  -- The monomorphic "return" operator
940                  -> DsM CoreExpr
941 dsInnerMonadComp stmts bndrs ret_op
942   = dsMcStmts (stmts ++ [noLoc (ReturnStmt (mkBigLHsVarTup bndrs) ret_op)])
943
944 -- The `unzip` function for `GroupStmt` in a monad comprehensions
945 --
946 --   unzip :: m (a,b,..) -> (m a,m b,..)
947 --   unzip m_tuple = ( liftM selN1 m_tuple
948 --                   , liftM selN2 m_tuple
949 --                   , .. )
950 --
951 --   mkMcUnzipM m [t1, t2]
952 --     = (unzip_fn, \ys :: m (t1, t2) ->
953 --         ( liftM (selN1 :: (t1, t2) -> t1) ys
954 --         , liftM (selN2 :: (t1, t2) -> t2) ys
955 --         ))
956 --
957 mkMcUnzipM :: CoreExpr
958            -> Type                      -- m
959            -> [Type]                    -- [a,b,c,..]
960            -> DsM (Id, CoreExpr)
961 mkMcUnzipM liftM_op m_ty elt_tys
962   = do  { ys    <- newSysLocalDs monad_tuple_ty
963         ; xs    <- mapM newSysLocalDs elt_tys
964         ; scrut <- newSysLocalDs tuple_tys
965
966         ; unzip_fn <- newSysLocalDs unzip_fn_ty
967
968         ; let -- Select one Id from our tuple
969               selectExpr n = mkLams [scrut] $ mkTupleSelector xs (xs !! n) scrut (Var scrut)
970               -- Apply 'selectVar' and 'ys' to 'liftM'
971               tupleElem n = mkApps liftM_op
972                                    -- Types (m is figured out by the type checker):
973                                    -- liftM :: forall a b. (a -> b) -> m a -> m b
974                                    [ Type tuple_tys, Type (elt_tys !! n)
975                                    -- Arguments:
976                                    , selectExpr n, Var ys ]
977               -- The final expression with the big tuple
978               unzip_body = mkBigCoreTup [ tupleElem n | n <- [0..length elt_tys - 1] ]
979
980         ; return (unzip_fn, mkLams [ys] unzip_body) }
981   where monad_tys       = map (m_ty `mkAppTy`) elt_tys                  -- [m a,m b,m c,..]
982         tuple_monad_tys = mkBigCoreTupTy monad_tys                      -- (m a,m b,m c,..)
983         tuple_tys       = mkBigCoreTupTy elt_tys                        -- (a,b,c,..)
984         monad_tuple_ty  = m_ty `mkAppTy` tuple_tys                      -- m (a,b,c,..)
985         unzip_fn_ty     = monad_tuple_ty `mkFunTy` tuple_monad_tys      -- m (a,b,c,..) -> (m a,m b,m c,..)
986
987 -- Generate the `mzip` function for `ParStmt` in monad comprehensions, for
988 -- example:
989 --
990 --   mzip :: m t1
991 --        -> (m t2 -> m t3 -> m (t2, t3))
992 --        -> m (t1, (t2, t3))
993 --
994 --   mkMcZipM m [t1, t2, t3]
995 --     = (zip_fn, \(q1::t1) (q2::t2) (q3::t3) ->
996 --         mzip q1 (mzip q2 q3))
997 --
998 mkMcZipM :: CoreExpr
999          -> Type
1000          -> [Type]
1001          -> DsM (Id, CoreExpr)
1002
1003 mkMcZipM mzip_op m_ty tys@(_:_:_) -- min. 2 types
1004  = do  { (ids, t1, tuple_ty, zip_body) <- loop tys
1005        ; zip_fn <- newSysLocalDs $
1006                        (m_ty `mkAppTy` t1)
1007                        `mkFunTy`
1008                        (m_ty `mkAppTy` tuple_ty)
1009                        `mkFunTy`
1010                        (m_ty `mkAppTy` mkBigCoreTupTy [t1, tuple_ty])
1011        ; return (zip_fn, mkLams ids zip_body) }
1012
1013  where 
1014        -- loop :: [Type] -> DsM ([Id], Type, [Type], CoreExpr)
1015        loop [t1, t2] = do -- last run of the `loop`
1016            { ids@[a,b] <- newSysLocalsDs (map (m_ty `mkAppTy`) [t1,t2])
1017            ; let zip_body = mkApps mzip_op [ Type t1, Type t2 , Var a, Var b ]
1018            ; return (ids, t1, t2, zip_body) }
1019
1020        loop (t1:tr) = do
1021            { -- Get ty, ids etc from the "inner" zip
1022              (ids', t1', t2', zip_body') <- loop tr
1023
1024            ; a <- newSysLocalDs $ m_ty `mkAppTy` t1
1025            ; let tuple_ty' = mkBigCoreTupTy [t1', t2']
1026                  zip_body = mkApps mzip_op [ Type t1, Type tuple_ty', Var a, zip_body' ]
1027            ; return ((a:ids'), t1, tuple_ty', zip_body) }
1028
1029 -- This case should never happen:
1030 mkMcZipM _ _ tys = pprPanic "mkMcZipM: unexpected argument" (ppr tys)
1031
1032 \end{code}