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 and array comprehensions
   7
   8 \begin{code}
   9 {-# OPTIONS -w #-}
  10 -- The above warning supression flag is a temporary kludge.
  11 -- While working on this module you are encouraged to remove it and fix
  12 -- any warnings in the module. See
  13 --     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
  14 -- for details
  15
  16 module DsListComp ( dsListComp, dsPArrComp ) where
  17
  18 #include "HsVersions.h"
  19
  20 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLocalBinds )
  21
  22 import BasicTypes
  23 import HsSyn
  24 import TcHsSyn
  25 import CoreSyn
  26
  27 import DsMonad          -- the monadery used in the desugarer
  28 import DsUtils
  29
  30 import DynFlags
  31 import CoreUtils
  32 import Var
  33 import Type
  34 import TysPrim
  35 import TysWiredIn
  36 import Match
  37 import PrelNames
  38 import PrelInfo
  39 import SrcLoc
  40 import Panic
  41 import Outputable
  42
  43 import Control.Monad ( liftM2 )
  44 \end{code}
  45
  46 List comprehensions may be desugared in one of two ways: ``ordinary''
  47 (as you would expect if you read SLPJ's book) and ``with foldr/build
  48 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
  49
  50 There will be at least one ``qualifier'' in the input.
  51
  52 \begin{code}
  53 dsListComp :: [LStmt Id]
  54            -> LHsExpr Id
  55            -> Type              -- Type of list elements
  56            -> DsM CoreExpr
  57 dsListComp lquals body elt_ty = do
  58     dflags <- getDOptsDs
  59     let quals = map unLoc lquals
  60
  61     if not (dopt Opt_RewriteRules 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 body (mkNilExpr elt_ty)
  68         else do -- Foldr/build should be enabled, so desugar
  69                 -- into foldrs and builds
  70             [n_tyvar] <- newTyVarsDs [alphaTyVar]
  71
  72             let n_ty = mkTyVarTy n_tyvar
  73                 c_ty = mkFunTys [elt_ty, n_ty] n_ty
  74             [c, n] <- newSysLocalsDs [c_ty, n_ty]
  75
  76             result <- dfListComp c n quals body
  77             build_id <- dsLookupGlobalId buildName
  78             returnDs (Var build_id `App` Type elt_ty `App` mkLams [n_tyvar, c, n] result)
  79
  80   where
  81     -- We must test for ParStmt anywhere, not just at the head, because an extension
  82     -- to list comprehensions would be to add brackets to specify the associativity
  83     -- of qualifier lists. This is really easy to do by adding extra ParStmts into the
  84     -- mix of possibly a single element in length, so we do this to leave the possibility open
  85     isParallelComp = any isParallelStmt
  86
  87     isParallelStmt (ParStmt _) = True
  88     isParallelStmt _           = False
  89
  90
  91 -- This function lets you desugar a inner list comprehension and a list of the binders
  92 -- of that comprehension that we need in the outer comprehension into such an expression
  93 -- and the type of the elements that it outputs (tuples of binders)
  94 dsInnerListComp :: ([LStmt Id], [Id]) -> DsM (CoreExpr, Type)
  95 dsInnerListComp (stmts, bndrs) = do
  96         expr <- dsListComp stmts (mkBigLHsVarTup bndrs) bndrs_tuple_type
  97         return (expr, bndrs_tuple_type)
  98     where
  99         bndrs_types = map idType bndrs
 100         bndrs_tuple_type = mkBigCoreTupTy bndrs_types
 101
 102
 103 -- This function factors out commonality between the desugaring strategies for TransformStmt.
 104 -- Given such a statement it gives you back an expression representing how to compute the transformed
 105 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
 106 dsTransformStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
 107 dsTransformStmt (TransformStmt (stmts, binders) usingExpr maybeByExpr) = do
 108     (expr, binders_tuple_type) <- dsInnerListComp (stmts, binders)
 109     usingExpr' <- dsLExpr usingExpr
 110
 111     using_args <-
 112         case maybeByExpr of
 113             Nothing -> return [expr]
 114             Just byExpr -> do
 115                 byExpr' <- dsLExpr byExpr
 116
 117                 us <- newUniqueSupply
 118                 [tuple_binder] <- newSysLocalsDs [binders_tuple_type]
 119                 let byExprWrapper = mkTupleCase us binders byExpr' tuple_binder (Var tuple_binder)
 120
 121                 return [Lam tuple_binder byExprWrapper, expr]
 122
 123     let inner_list_expr = mkApps usingExpr' ((Type binders_tuple_type) : using_args)
 124
 125     let pat = mkBigLHsVarPatTup binders
 126     return (inner_list_expr, pat)
 127
 128 -- This function factors out commonality between the desugaring strategies for GroupStmt.
 129 -- Given such a statement it gives you back an expression representing how to compute the transformed
 130 -- list and the tuple that you need to bind from that list in order to proceed with your desugaring
 131 dsGroupStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
 132 dsGroupStmt (GroupStmt (stmts, binderMap) groupByClause) = do
 133     let (fromBinders, toBinders) = unzip binderMap
 134
 135         fromBindersTypes = map idType fromBinders
 136         toBindersTypes = map idType toBinders
 137
 138         toBindersTupleType = mkBigCoreTupTy toBindersTypes
 139
 140     -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
 141     (expr, fromBindersTupleType) <- dsInnerListComp (stmts, fromBinders)
 142
 143     -- Work out what arguments should be supplied to that expression: i.e. is an extraction
 144     -- function required? If so, create that desugared function and add to arguments
 145     (usingExpr', usingArgs) <-
 146         case groupByClause of
 147             GroupByNothing usingExpr -> liftM2 (,) (dsLExpr usingExpr) (return [expr])
 148             GroupBySomething usingExpr byExpr -> do
 149                 usingExpr' <- dsLExpr (either id noLoc usingExpr)
 150
 151                 byExpr' <- dsLExpr byExpr
 152
 153                 us <- newUniqueSupply
 154                 [fromBindersTuple] <- newSysLocalsDs [fromBindersTupleType]
 155                 let byExprWrapper = mkTupleCase us fromBinders byExpr' fromBindersTuple (Var fromBindersTuple)
 156
 157                 return (usingExpr', [Lam fromBindersTuple byExprWrapper, expr])
 158
 159     -- Create an unzip function for the appropriate arity and element types and find "map"
 160     (unzip_fn, unzip_rhs) <- mkUnzipBind fromBindersTypes
 161     map_id <- dsLookupGlobalId mapName
 162
 163     -- Generate the expressions to build the grouped list
 164     let -- First we apply the grouping function to the inner list
 165         inner_list_expr = mkApps usingExpr' ((Type fromBindersTupleType) : usingArgs)
 166         -- Then we map our "unzip" across it to turn the lists of tuples into tuples of lists
 167         -- We make sure we instantiate the type variable "a" to be a list of "from" tuples and
 168         -- the "b" to be a tuple of "to" lists!
 169         unzipped_inner_list_expr = mkApps (Var map_id)
 170             [Type (mkListTy fromBindersTupleType), Type toBindersTupleType, Var unzip_fn, inner_list_expr]
 171         -- Then finally we bind the unzip function around that expression
 172         bound_unzipped_inner_list_expr = Let (Rec [(unzip_fn, unzip_rhs)]) unzipped_inner_list_expr
 173
 174     -- Build a pattern that ensures the consumer binds into the NEW binders, which hold lists rather than single values
 175     let pat = mkBigLHsVarPatTup toBinders
 176     return (bound_unzipped_inner_list_expr, pat)
 177
 178 \end{code}
 179
 180 %************************************************************************
 181 %*                                                                      *
 182 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
 183 %*                                                                      *
 184 %************************************************************************
 185
 186 Just as in Phil's chapter~7 in SLPJ, using the rules for
 187 optimally-compiled list comprehensions.  This is what Kevin followed
 188 as well, and I quite happily do the same.  The TQ translation scheme
 189 transforms a list of qualifiers (either boolean expressions or
 190 generators) into a single expression which implements the list
 191 comprehension.  Because we are generating 2nd-order polymorphic
 192 lambda-calculus, calls to NIL and CONS must be applied to a type
 193 argument, as well as their usual value arguments.
 194 \begin{verbatim}
 195 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
 196
 197 (Rule C)
 198 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
 199
 200 (Rule B)
 201 TQ << [ e | b , qs ] ++ L >> =
 202     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
 203
 204 (Rule A')
 205 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
 206   letrec
 207     h = \ u1 ->
 208           case u1 of
 209             []        ->  TE << L2 >>
 210             (u2 : u3) ->
 211                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
 212                     [] (h u3)
 213   in
 214     h ( TE << L1 >> )
 215
 216 "h", "u1", "u2", and "u3" are new variables.
 217 \end{verbatim}
 218
 219 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
 220 is the TE translation scheme.  Note that we carry around the @L@ list
 221 already desugared.  @dsListComp@ does the top TE rule mentioned above.
 222
 223 To the above, we add an additional rule to deal with parallel list
 224 comprehensions.  The translation goes roughly as follows:
 225      [ e | p1 <- e11, let v1 = e12, p2 <- e13
 226          | q1 <- e21, let v2 = e22, q2 <- e23]
 227      =>
 228      [ e | ((x1, .., xn), (y1, ..., ym)) <-
 229                zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
 230                    [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
 231 where (x1, .., xn) are the variables bound in p1, v1, p2
 232       (y1, .., ym) are the variables bound in q1, v2, q2
 233
 234 In the translation below, the ParStmt branch translates each parallel branch
 235 into a sub-comprehension, and desugars each independently.  The resulting lists
 236 are fed to a zip function, we create a binding for all the variables bound in all
 237 the comprehensions, and then we hand things off the the desugarer for bindings.
 238 The zip function is generated here a) because it's small, and b) because then we
 239 don't have to deal with arbitrary limits on the number of zip functions in the
 240 prelude, nor which library the zip function came from.
 241 The introduced tuples are Boxed, but only because I couldn't get it to work
 242 with the Unboxed variety.
 243
 244 \begin{code}
 245
 246 deListComp :: [Stmt Id] -> LHsExpr Id -> CoreExpr -> DsM CoreExpr
 247
 248 deListComp (ParStmt stmtss_w_bndrs : quals) body list
 249   = do
 250     exps_and_qual_tys <- mappM dsInnerListComp stmtss_w_bndrs
 251     let (exps, qual_tys) = unzip exps_and_qual_tys
 252
 253     (zip_fn, zip_rhs) <- mkZipBind qual_tys
 254
 255         -- Deal with [e | pat <- zip l1 .. ln] in example above
 256     deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
 257                    quals body list
 258
 259   where
 260         bndrs_s = map snd stmtss_w_bndrs
 261
 262         -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
 263         pat      = mkBigLHsPatTup pats
 264         pats = map mkBigLHsVarPatTup bndrs_s
 265
 266         -- Last: the one to return
 267 deListComp [] body list         -- Figure 7.4, SLPJ, p 135, rule C above
 268   = dsLExpr body                `thenDs` \ core_body ->
 269     returnDs (mkConsExpr (exprType core_body) core_body list)
 270
 271         -- Non-last: must be a guard
 272 deListComp (ExprStmt guard _ _ : quals) body list       -- rule B above
 273   = dsLExpr guard               `thenDs` \ core_guard ->
 274     deListComp quals body list  `thenDs` \ core_rest ->
 275     returnDs (mkIfThenElse core_guard core_rest list)
 276
 277 -- [e | let B, qs] = let B in [e | qs]
 278 deListComp (LetStmt binds : quals) body list
 279   = deListComp quals body list  `thenDs` \ core_rest ->
 280     dsLocalBinds binds core_rest
 281
 282 deListComp (stmt@(TransformStmt _ _ _) : quals) body list = do
 283     (inner_list_expr, pat) <- dsTransformStmt stmt
 284     deBindComp pat inner_list_expr quals body list
 285
 286 deListComp (stmt@(GroupStmt _ _) : quals) body list = do
 287     (inner_list_expr, pat) <- dsGroupStmt stmt
 288     deBindComp pat inner_list_expr quals body list
 289
 290 deListComp (BindStmt pat list1 _ _ : quals) body core_list2 -- rule A' above
 291   = dsLExpr list1                   `thenDs` \ core_list1 ->
 292     deBindComp pat core_list1 quals body core_list2
 293 \end{code}
 294
 295
 296 \begin{code}
 297 deBindComp pat core_list1 quals body core_list2 = do
 298     let
 299         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
 300
 301         -- u1_ty is a [alpha] type, and u2_ty = alpha
 302         u2_ty = hsLPatType pat
 303
 304         res_ty = exprType core_list2
 305         h_ty   = u1_ty `mkFunTy` res_ty
 306
 307     [h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]
 308
 309     -- the "fail" value ...
 310     let
 311         core_fail   = App (Var h) (Var u3)
 312         letrec_body = App (Var h) core_list1
 313
 314     rest_expr <- deListComp quals body core_fail
 315     core_match <- matchSimply (Var u2) (StmtCtxt ListComp) pat rest_expr core_fail
 316
 317     let
 318         rhs = Lam u1 $
 319               Case (Var u1) u1 res_ty
 320                    [(DataAlt nilDataCon,  [],       core_list2),
 321                     (DataAlt consDataCon, [u2, u3], core_match)]
 322                         -- Increasing order of tag
 323
 324     return (Let (Rec [(h, rhs)]) letrec_body)
 325 \end{code}
 326
 327 %************************************************************************
 328 %*                                                                      *
 329 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 330 %*                                                                      *
 331 %************************************************************************
 332
 333 @dfListComp@ are the rules used with foldr/build turned on:
 334
 335 \begin{verbatim}
 336 TE[ e | ]            c n = c e n
 337 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
 338 TE[ e | p <- l , q ] c n = let
 339                                 f = \ x b -> case x of
 340                                                   p -> TE[ e | q ] c b
 341                                                   _ -> b
 342                            in
 343                            foldr f n l
 344 \end{verbatim}
 345
 346 \begin{code}
 347 dfListComp :: Id -> Id -- 'c' and 'n'
 348         -> [Stmt Id]   -- the rest of the qual's
 349         -> LHsExpr Id
 350         -> DsM CoreExpr
 351
 352         -- Last: the one to return
 353 dfListComp c_id n_id [] body
 354   = dsLExpr body                `thenDs` \ core_body ->
 355     returnDs (mkApps (Var c_id) [core_body, Var n_id])
 356
 357         -- Non-last: must be a guard
 358 dfListComp c_id n_id (ExprStmt guard _ _  : quals) body
 359   = dsLExpr guard                       `thenDs` \ core_guard ->
 360     dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
 361     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 362
 363 dfListComp c_id n_id (LetStmt binds : quals) body
 364   -- new in 1.3, local bindings
 365   = dfListComp c_id n_id quals body     `thenDs` \ core_rest ->
 366     dsLocalBinds binds core_rest
 367
 368 dfListComp c_id n_id (stmt@(TransformStmt _ _ _) : quals) body = do
 369     (inner_list_expr, pat) <- dsTransformStmt stmt
 370     -- Anyway, we bind the newly transformed list via the generic binding function
 371     dfBindComp c_id n_id (pat, inner_list_expr) quals body
 372
 373 dfListComp c_id n_id (stmt@(GroupStmt _ _) : quals) body = do
 374     (inner_list_expr, pat) <- dsGroupStmt stmt
 375     -- Anyway, we bind the newly grouped list via the generic binding function
 376     dfBindComp c_id n_id (pat, inner_list_expr) quals body
 377
 378 dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) body = do
 379     -- evaluate the two lists
 380     core_list1 <- dsLExpr list1
 381
 382     -- Do the rest of the work in the generic binding builder
 383     dfBindComp c_id n_id (pat, core_list1) quals body
 384
 385 dfBindComp :: Id -> Id          -- 'c' and 'n'
 386        -> (LPat Id, CoreExpr)
 387            -> [Stmt Id]                 -- the rest of the qual's
 388            -> LHsExpr Id
 389            -> DsM CoreExpr
 390 dfBindComp c_id n_id (pat, core_list1) quals body = do
 391     -- find the required type
 392     let x_ty   = hsLPatType pat
 393         b_ty   = idType n_id
 394
 395     -- create some new local id's
 396     [b, x] <- newSysLocalsDs [b_ty, x_ty]
 397
 398     -- build rest of the comprehesion
 399     core_rest <- dfListComp c_id b quals body
 400
 401     -- build the pattern match
 402     core_expr <- matchSimply (Var x) (StmtCtxt ListComp)
 403                 pat core_rest (Var b)
 404
 405     -- now build the outermost foldr, and return
 406     foldr_id <- dsLookupGlobalId foldrName
 407     return (Var foldr_id `App` Type x_ty
 408                `App` Type b_ty
 409                `App` mkLams [x, b] core_expr
 410                `App` Var n_id
 411                `App` core_list1)
 412
 413 \end{code}
 414
 415 %************************************************************************
 416 %*                                                                      *
 417 \subsection[DsFunGeneration]{Generation of zip/unzip functions for use in desugaring}
 418 %*                                                                      *
 419 %************************************************************************
 420
 421 \begin{code}
 422
 423 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
 424 -- mkZipBind [t1, t2]
 425 -- = (zip, \as1:[t1] as2:[t2]
 426 --         -> case as1 of
 427 --              [] -> []
 428 --              (a1:as'1) -> case as2 of
 429 --                              [] -> []
 430 --                              (a2:as'2) -> (a1, a2) : zip as'1 as'2)]
 431
 432 mkZipBind elt_tys = do
 433     ass  <- mappM newSysLocalDs  elt_list_tys
 434     as'  <- mappM newSysLocalDs  elt_tys
 435     as's <- mappM newSysLocalDs  elt_list_tys
 436
 437     zip_fn <- newSysLocalDs zip_fn_ty
 438
 439     let inner_rhs = mkConsExpr elt_tuple_ty
 440                         (mkBigCoreVarTup as')
 441                         (mkVarApps (Var zip_fn) as's)
 442         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
 443
 444     return (zip_fn, mkLams ass zip_body)
 445   where
 446     elt_list_tys      = map mkListTy elt_tys
 447     elt_tuple_ty      = mkBigCoreTupTy elt_tys
 448     elt_tuple_list_ty = mkListTy elt_tuple_ty
 449
 450     zip_fn_ty         = mkFunTys elt_list_tys elt_tuple_list_ty
 451
 452     mk_case (as, a', as') rest
 453           = Case (Var as) as elt_tuple_list_ty
 454                   [(DataAlt nilDataCon,  [],        mkNilExpr elt_tuple_ty),
 455                    (DataAlt consDataCon, [a', as'], rest)]
 456                         -- Increasing order of tag
 457
 458
 459 mkUnzipBind :: [Type] -> DsM (Id, CoreExpr)
 460 -- mkUnzipBind [t1, t2]
 461 -- = (unzip, \ys :: [(t1, t2)] -> foldr (\ax :: (t1, t2) axs :: ([t1], [t2])
 462 --     -> case ax of
 463 --      (x1, x2) -> case axs of
 464 --                (xs1, xs2) -> (x1 : xs1, x2 : xs2))
 465 --      ([], [])
 466 --      ys)
 467 --
 468 -- We use foldr here in all cases, even if rules are turned off, because we may as well!
 469 mkUnzipBind elt_tys = do
 470     ax  <- newSysLocalDs elt_tuple_ty
 471     axs <- newSysLocalDs elt_list_tuple_ty
 472     ys  <- newSysLocalDs elt_tuple_list_ty
 473     xs  <- mappM newSysLocalDs elt_tys
 474     xss <- mappM newSysLocalDs elt_list_tys
 475
 476     unzip_fn <- newSysLocalDs unzip_fn_ty
 477
 478     foldr_id <- dsLookupGlobalId foldrName
 479     [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
 480
 481     let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
 482
 483         concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
 484         tupled_concat_expression = mkBigCoreTup concat_expressions
 485
 486         folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
 487         folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
 488         folder_body = mkLams [ax, axs] folder_body_outer_case
 489
 490         unzip_body = mkApps (Var foldr_id) [Type elt_tuple_ty, Type elt_list_tuple_ty, folder_body, nil_tuple, Var ys]
 491         unzip_body_saturated = mkLams [ys] unzip_body
 492
 493     return (unzip_fn, unzip_body_saturated)
 494   where
 495     elt_tuple_ty       = mkBigCoreTupTy elt_tys
 496     elt_tuple_list_ty  = mkListTy elt_tuple_ty
 497     elt_list_tys       = map mkListTy elt_tys
 498     elt_list_tuple_ty  = mkBigCoreTupTy elt_list_tys
 499
 500     unzip_fn_ty        = elt_tuple_list_ty `mkFunTy` elt_list_tuple_ty
 501
 502     mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
 503
 504
 505
 506 \end{code}
 507
 508 %************************************************************************
 509 %*                                                                      *
 510 \subsection[DsPArrComp]{Desugaring of array comprehensions}
 511 %*                                                                      *
 512 %************************************************************************
 513
 514 \begin{code}
 515
 516 -- entry point for desugaring a parallel array comprehension
 517 --
 518 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
 519 --
 520 dsPArrComp :: [Stmt Id]
 521             -> LHsExpr Id
 522             -> Type                 -- Don't use; called with `undefined' below
 523             -> DsM CoreExpr
 524 dsPArrComp [ParStmt qss] body _  =  -- parallel comprehension
 525   dePArrParComp qss body
 526 dsPArrComp qs            body _  =  -- no ParStmt in `qs'
 527   dsLookupGlobalId singletonPName                         `thenDs` \sglP ->
 528   let unitArray = mkApps (Var sglP) [Type unitTy,
 529                                      mkCoreTup []]
 530   in
 531   dePArrComp qs body (mkLHsPatTup []) unitArray
 532
 533
 534
 535 -- the work horse
 536 --
 537 dePArrComp :: [Stmt Id]
 538            -> LHsExpr Id
 539            -> LPat Id           -- the current generator pattern
 540            -> CoreExpr          -- the current generator expression
 541            -> DsM CoreExpr
 542 --
 543 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
 544 --
 545 dePArrComp [] e' pa cea =
 546   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 547   let ty = parrElemType cea
 548   in
 549   deLambda ty pa e'                                       `thenDs` \(clam,
 550                                                                      ty'e') ->
 551   returnDs $ 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) body pa cea =
 556   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
 557   let ty = parrElemType cea
 558   in
 559   deLambda ty pa b                                `thenDs` \(clam,_) ->
 560   dePArrComp qs body pa (mkApps (Var filterP) [Type ty, clam, cea])
 561
 562 --
 563 --  <<[:e' | p <- e, qs:]>> pa ea =
 564 --    let ef = \pa -> e
 565 --    in
 566 --    <<[:e' | qs:]>> (pa, p) (crossMap ea ef)
 567 --
 568 -- if matching again p cannot fail, or else
 569 --
 570 --  <<[:e' | p <- e, qs:]>> pa ea =
 571 --    let ef = \pa -> filterP (\x -> case x of {p -> True; _ -> False}) e
 572 --    in
 573 --    <<[:e' | qs:]>> (pa, p) (crossMapP ea ef)
 574 --
 575 dePArrComp (BindStmt p e _ _ : qs) body pa cea =
 576   dsLookupGlobalId filterPName                    `thenDs` \filterP    ->
 577   dsLookupGlobalId crossMapPName                  `thenDs` \crossMapP  ->
 578   dsLExpr e                                       `thenDs` \ce         ->
 579   let ety'cea = parrElemType cea
 580       ety'ce  = parrElemType ce
 581       false   = Var falseDataConId
 582       true    = Var trueDataConId
 583   in
 584   newSysLocalDs ety'ce                                    `thenDs` \v       ->
 585   matchSimply (Var v) (StmtCtxt PArrComp) p true false    `thenDs` \pred    ->
 586   let cef | isIrrefutableHsPat p = ce
 587           | otherwise            = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
 588   in
 589   mkLambda ety'cea pa cef                                 `thenDs` \(clam,
 590                                                                      _    ) ->
 591   let ety'cef = ety'ce              -- filter doesn't change the element type
 592       pa'     = mkLHsPatTup [pa, p]
 593   in
 594   dePArrComp qs body pa' (mkApps (Var crossMapP)
 595                                  [Type ety'cea, Type ety'cef, cea, clam])
 596 --
 597 --  <<[:e' | let ds, qs:]>> pa ea =
 598 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
 599 --                    (mapP (\v@pa -> let ds in (v, (x_1, ..., x_n))) ea)
 600 --  where
 601 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
 602 --
 603 dePArrComp (LetStmt ds : qs) body pa cea =
 604   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 605   let xs     = map unLoc (collectLocalBinders ds)
 606       ty'cea = parrElemType cea
 607   in
 608   newSysLocalDs ty'cea                                    `thenDs` \v       ->
 609   dsLocalBinds ds (mkCoreTup (map Var xs))                `thenDs` \clet    ->
 610   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
 611   let projBody = mkDsLet (NonRec let'v clet) $
 612                  mkCoreTup [Var v, Var let'v]
 613       errTy    = exprType projBody
 614       errMsg   = "DsListComp.dePArrComp: internal error!"
 615   in
 616   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 617   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr`thenDs` \ccase   ->
 618   let pa'    = mkLHsPatTup [pa, mkLHsPatTup (map nlVarPat xs)]
 619       proj   = mkLams [v] ccase
 620   in
 621   dePArrComp qs body pa' (mkApps (Var mapP)
 622                                  [Type ty'cea, Type errTy, proj, cea])
 623 --
 624 -- The parser guarantees that parallel comprehensions can only appear as
 625 -- singeltons qualifier lists, which we already special case in the caller.
 626 -- So, encountering one here is a bug.
 627 --
 628 dePArrComp (ParStmt _ : _) _ _ _ =
 629   panic "DsListComp.dePArrComp: malformed comprehension AST"
 630
 631 --  <<[:e' | qs | qss:]>> pa ea =
 632 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
 633 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
 634 --    where
 635 --      {x_1, ..., x_n} = DV (qs)
 636 --
 637 dePArrParComp qss body =
 638   deParStmt qss                                         `thenDs` \(pQss,
 639                                                                    ceQss) ->
 640   dePArrComp [] body pQss ceQss
 641   where
 642     deParStmt []             =
 643       -- empty parallel statement lists have no source representation
 644       panic "DsListComp.dePArrComp: Empty parallel list comprehension"
 645     deParStmt ((qs, xs):qss) =          -- first statement
 646       let res_expr = mkLHsVarTup xs
 647       in
 648       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
 649       parStmts qss (mkLHsVarPatTup xs) cqs
 650     ---
 651     parStmts []             pa cea = return (pa, cea)
 652     parStmts ((qs, xs):qss) pa cea =    -- subsequent statements (zip'ed)
 653       dsLookupGlobalId zipPName                           `thenDs` \zipP    ->
 654       let pa'      = mkLHsPatTup [pa, mkLHsVarPatTup xs]
 655           ty'cea   = parrElemType cea
 656           res_expr = mkLHsVarTup xs
 657       in
 658       dsPArrComp (map unLoc qs) res_expr undefined        `thenDs` \cqs     ->
 659       let ty'cqs = parrElemType cqs
 660           cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
 661       in
 662       parStmts qss pa' cea'
 663
 664 -- generate Core corresponding to `\p -> e'
 665 --
 666 deLambda :: Type                        -- type of the argument
 667           -> LPat Id                    -- argument pattern
 668           -> LHsExpr Id                 -- body
 669           -> DsM (CoreExpr, Type)
 670 deLambda ty p e =
 671   dsLExpr e                                               `thenDs` \ce      ->
 672   mkLambda ty p ce
 673
 674 -- generate Core for a lambda pattern match, where the body is already in Core
 675 --
 676 mkLambda :: Type                        -- type of the argument
 677          -> LPat Id                     -- argument pattern
 678          -> CoreExpr                    -- desugared body
 679          -> DsM (CoreExpr, Type)
 680 mkLambda ty p ce =
 681   newSysLocalDs ty                                        `thenDs` \v       ->
 682   let errMsg = "DsListComp.deLambda: internal error!"
 683       ce'ty  = exprType ce
 684   in
 685   mkErrorAppDs pAT_ERROR_ID ce'ty errMsg                  `thenDs` \cerr    ->
 686   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
 687   returnDs (mkLams [v] res, ce'ty)
 688
 689 -- obtain the element type of the parallel array produced by the given Core
 690 -- expression
 691 --
 692 parrElemType   :: CoreExpr -> Type
 693 parrElemType e  =
 694   case splitTyConApp_maybe (exprType e) of
 695     Just (tycon, [ty]) | tycon == parrTyCon -> ty
 696     _                                                     -> panic
 697       "DsListComp.parrElemType: not a parallel array type"
 698 \end{code}