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