ghc/compiler/deSugar/DsListComp.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[DsListComp]{Desugaring list comprehensions and array comprehensions}
   5
   6 \begin{code}
   7 module DsListComp ( dsListComp, dsPArrComp ) where
   8
   9 #include "HsVersions.h"
  10
  11 import {-# SOURCE #-} DsExpr ( dsLExpr, dsLet )
  12
  13 import BasicTypes       ( Boxity(..) )
  14 import HsSyn
  15 import TcHsSyn          ( hsPatType )
  16 import CoreSyn
  17
  18 import DsMonad          -- the monadery used in the desugarer
  19 import DsUtils
  20
  21 import CmdLineOpts      ( DynFlag(..), dopt, opt_RulesOff )
  22 import CoreUtils        ( exprType, mkIfThenElse )
  23 import Id               ( idType )
  24 import Var              ( Id )
  25 import Type             ( mkTyVarTy, mkFunTys, mkFunTy, Type,
  26                           splitTyConApp_maybe )
  27 import TysPrim          ( alphaTyVar )
  28 import TysWiredIn       ( nilDataCon, consDataCon, trueDataConId, falseDataConId,
  29                           unitDataConId, unitTy, mkListTy, parrTyCon )
  30 import Match            ( matchSimply )
  31 import PrelNames        ( foldrName, buildName, replicatePName, mapPName,
  32                           filterPName, zipPName, crossPName )
  33 import PrelInfo         ( pAT_ERROR_ID )
  34 import SrcLoc           ( noLoc, unLoc )
  35 import Panic            ( panic )
  36 \end{code}
  37
  38 List comprehensions may be desugared in one of two ways: ``ordinary''
  39 (as you would expect if you read SLPJ's book) and ``with foldr/build
  40 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
  41
  42 There will be at least one ``qualifier'' in the input.
  43
  44 \begin{code}
  45 dsListComp :: [LStmt Id]
  46            -> Type              -- Type of list elements
  47            -> DsM CoreExpr
  48 dsListComp lquals elt_ty
  49   = getDOptsDs  `thenDs` \dflags ->
  50     let
  51         quals = map unLoc lquals
  52     in
  53     if opt_RulesOff || dopt Opt_IgnoreInterfacePragmas dflags
  54         -- Either rules are switched off, or we are ignoring what there are;
  55         -- Either way foldr/build won't happen, so use the more efficient
  56         -- Wadler-style desugaring
  57         || isParallelComp quals
  58                 -- Foldr-style desugaring can't handle
  59                 -- parallel list comprehensions
  60         then deListComp quals (mkNilExpr elt_ty)
  61
  62    else         -- Foldr/build should be enabled, so desugar
  63                 -- into foldrs and builds
  64     newTyVarsDs [alphaTyVar]    `thenDs` \ [n_tyvar] ->
  65     let
  66         n_ty = mkTyVarTy n_tyvar
  67         c_ty = mkFunTys [elt_ty, n_ty] n_ty
  68     in
  69     newSysLocalsDs [c_ty,n_ty]          `thenDs` \ [c, n] ->
  70     dfListComp c n quals                `thenDs` \ result ->
  71     dsLookupGlobalId buildName  `thenDs` \ build_id ->
  72     returnDs (Var build_id `App` Type elt_ty
  73                            `App` mkLams [n_tyvar, c, n] result)
  74
  75   where isParallelComp (ParStmt bndrstmtss : _) = True
  76         isParallelComp _                        = False
  77 \end{code}
  78
  79 %************************************************************************
  80 %*                                                                      *
  81 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
  82 %*                                                                      *
  83 %************************************************************************
  84
  85 Just as in Phil's chapter~7 in SLPJ, using the rules for
  86 optimally-compiled list comprehensions.  This is what Kevin followed
  87 as well, and I quite happily do the same.  The TQ translation scheme
  88 transforms a list of qualifiers (either boolean expressions or
  89 generators) into a single expression which implements the list
  90 comprehension.  Because we are generating 2nd-order polymorphic
  91 lambda-calculus, calls to NIL and CONS must be applied to a type
  92 argument, as well as their usual value arguments.
  93 \begin{verbatim}
  94 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
  95
  96 (Rule C)
  97 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
  98
  99 (Rule B)
 100 TQ << [ e | b , qs ] ++ L >> =
 101     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
 102
 103 (Rule A')
 104 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
 105   letrec
 106     h = \ u1 ->
 107           case u1 of
 108             []        ->  TE << L2 >>
 109             (u2 : u3) ->
 110                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
 111                     [] (h u3)
 112   in
 113     h ( TE << L1 >> )
 114
 115 "h", "u1", "u2", and "u3" are new variables.
 116 \end{verbatim}
 117
 118 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
 119 is the TE translation scheme.  Note that we carry around the @L@ list
 120 already desugared.  @dsListComp@ does the top TE rule mentioned above.
 121
 122 To the above, we add an additional rule to deal with parallel list
 123 comprehensions.  The translation goes roughly as follows:
 124      [ e | p1 <- e11, let v1 = e12, p2 <- e13
 125          | q1 <- e21, let v2 = e22, q2 <- e23]
 126      =>
 127      [ e | ((x1, .., xn), (y1, ..., ym)) <-
 128                zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
 129                    [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
 130 where (x1, .., xn) are the variables bound in p1, v1, p2
 131       (y1, .., ym) are the variables bound in q1, v2, q2
 132
 133 In the translation below, the ParStmt branch translates each parallel branch
 134 into a sub-comprehension, and desugars each independently.  The resulting lists
 135 are fed to a zip function, we create a binding for all the variables bound in all
 136 the comprehensions, and then we hand things off the the desugarer for bindings.
 137 The zip function is generated here a) because it's small, and b) because then we
 138 don't have to deal with arbitrary limits on the number of zip functions in the
 139 prelude, nor which library the zip function came from.
 140 The introduced tuples are Boxed, but only because I couldn't get it to work
 141 with the Unboxed variety.
 142
 143 \begin{code}
 144 deListComp :: [Stmt Id] -> CoreExpr -> DsM CoreExpr
 145
 146 deListComp (ParStmt stmtss_w_bndrs : quals) list
 147   = mappM do_list_comp stmtss_w_bndrs   `thenDs` \ exps ->
 148     mkZipBind qual_tys                  `thenDs` \ (zip_fn, zip_rhs) ->
 149
 150         -- Deal with [e | pat <- zip l1 .. ln] in example above
 151     deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
 152                    quals list
 153
 154   where
 155         bndrs_s = map snd stmtss_w_bndrs
 156
 157         -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
 158         pat      = noLoc (TuplePat pats Boxed)
 159         pats     = map mk_hs_tuple_pat bndrs_s
 160
 161         -- Types of (x1,..,xn), (y1,..,yn) etc
 162         qual_tys = map mk_bndrs_tys bndrs_s
 163
 164         do_list_comp (stmts, bndrs)
 165           = dsListComp (stmts ++ [noLoc $ ResultStmt (mk_hs_tuple_expr bndrs)])
 166                        (mk_bndrs_tys bndrs)
 167
 168         mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
 169
 170         -- Last: the one to return
 171 deListComp [ResultStmt expr] list       -- Figure 7.4, SLPJ, p 135, rule C above
 172   = dsLExpr expr                `thenDs` \ core_expr ->
 173     returnDs (mkConsExpr (exprType core_expr) core_expr list)
 174
 175         -- Non-last: must be a guard
 176 deListComp (ExprStmt guard ty : quals) list     -- rule B above
 177   = dsLExpr guard               `thenDs` \ core_guard ->
 178     deListComp quals list       `thenDs` \ core_rest ->
 179     returnDs (mkIfThenElse core_guard core_rest list)
 180
 181 -- [e | let B, qs] = let B in [e | qs]
 182 deListComp (LetStmt binds : quals) list
 183   = deListComp quals list       `thenDs` \ core_rest ->
 184     dsLet binds core_rest
 185
 186 deListComp (BindStmt pat list1 : quals) core_list2 -- rule A' above
 187   = dsLExpr list1                   `thenDs` \ core_list1 ->
 188     deBindComp pat core_list1 quals core_list2
 189 \end{code}
 190
 191
 192 \begin{code}
 193 deBindComp pat core_list1 quals core_list2
 194   = let
 195         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
 196
 197         -- u1_ty is a [alpha] type, and u2_ty = alpha
 198         u2_ty = hsPatType pat
 199
 200         res_ty = exprType core_list2
 201         h_ty   = u1_ty `mkFunTy` res_ty
 202     in
 203     newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]  `thenDs` \ [h, u1, u2, u3] ->
 204
 205     -- the "fail" value ...
 206     let
 207         core_fail   = App (Var h) (Var u3)
 208         letrec_body = App (Var h) core_list1
 209     in
 210     deListComp quals core_fail                  `thenDs` \ rest_expr ->
 211     matchSimply (Var u2) (StmtCtxt ListComp) pat
 212                 rest_expr core_fail             `thenDs` \ core_match ->
 213     let
 214         rhs = Lam u1 $
 215 -- gaw 2004
 216               Case (Var u1) u1 res_ty
 217                    [(DataAlt nilDataCon,  [],       core_list2),
 218                     (DataAlt consDataCon, [u2, u3], core_match)]
 219     in
 220     returnDs (Let (Rec [(h, rhs)]) letrec_body)
 221 \end{code}
 222
 223
 224 \begin{code}
 225 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
 226 -- mkZipBind [t1, t2]
 227 -- = (zip, \as1:[t1] as2:[t2]
 228 --         -> case as1 of
 229 --              [] -> []
 230 --              (a1:as'1) -> case as2 of
 231 --                              [] -> []
 232 --                              (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
 233
 234 mkZipBind elt_tys
 235   = mappM newSysLocalDs  list_tys       `thenDs` \ ass ->
 236     mappM newSysLocalDs  elt_tys        `thenDs` \ as' ->
 237     mappM newSysLocalDs  list_tys       `thenDs` \ as's ->
 238     newSysLocalDs zip_fn_ty             `thenDs` \ zip_fn ->
 239     let
 240         inner_rhs = mkConsExpr ret_elt_ty
 241                         (mkCoreTup (map Var as'))
 242                         (mkVarApps (Var zip_fn) as's)
 243         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
 244     in
 245     returnDs (zip_fn, mkLams ass zip_body)
 246   where
 247     list_tys    = map mkListTy elt_tys
 248     ret_elt_ty  = mkCoreTupTy elt_tys
 249     list_ret_ty = mkListTy ret_elt_ty
 250     zip_fn_ty   = mkFunTys list_tys list_ret_ty
 251
 252     mk_case (as, a', as') rest
 253 -- gaw 2004
 254           = Case (Var as) as list_ret_ty
 255                   [(DataAlt nilDataCon,  [],        mkNilExpr ret_elt_ty),
 256                    (DataAlt consDataCon, [a', as'], rest)]
 257
 258 -- Helper functions that makes an HsTuple only for non-1-sized tuples
 259 mk_hs_tuple_expr :: [Id] -> LHsExpr Id
 260 mk_hs_tuple_expr []   = nlHsVar unitDataConId
 261 mk_hs_tuple_expr [id] = nlHsVar id
 262 mk_hs_tuple_expr ids  = noLoc $ ExplicitTuple [ nlHsVar i | i <- ids ] Boxed
 263
 264 mk_hs_tuple_pat :: [Id] -> LPat Id
 265 mk_hs_tuple_pat [b] = nlVarPat b
 266 mk_hs_tuple_pat bs  = noLoc $ TuplePat (map nlVarPat bs) Boxed
 267 \end{code}
 268
 269
 270 %************************************************************************
 271 %*                                                                      *
 272 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 273 %*                                                                      *
 274 %************************************************************************
 275
 276 @dfListComp@ are the rules used with foldr/build turned on:
 277
 278 \begin{verbatim}
 279 TE[ e | ]            c n = c e n
 280 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
 281 TE[ e | p <- l , q ] c n = let
 282                                 f = \ x b -> case x of
 283                                                   p -> TE[ e | q ] c b
 284                                                   _ -> b
 285                            in
 286                            foldr f n l
 287 \end{verbatim}
 288
 289 \begin{code}
 290 dfListComp :: Id -> Id                  -- 'c' and 'n'
 291            -> [Stmt Id]         -- the rest of the qual's
 292            -> DsM CoreExpr
 293
 294         -- Last: the one to return
 295 dfListComp c_id n_id [ResultStmt expr]
 296   = dsLExpr expr                        `thenDs` \ core_expr ->
 297     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 298
 299         -- Non-last: must be a guard
 300 dfListComp c_id n_id (ExprStmt guard ty  : quals)
 301   = dsLExpr guard                               `thenDs` \ core_guard ->
 302     dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 303     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 304
 305 dfListComp c_id n_id (LetStmt binds : quals)
 306   -- new in 1.3, local bindings
 307   = dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 308     dsLet binds core_rest
 309
 310 dfListComp c_id n_id (BindStmt pat list1 : quals)
 311     -- evaluate the two lists
 312   = dsLExpr list1                       `thenDs` \ core_list1 ->
 313
 314     -- find the required type
 315     let x_ty   = hsPatType pat
 316         b_ty   = idType n_id
 317     in
 318
 319     -- create some new local id's
 320     newSysLocalsDs [b_ty,x_ty]                  `thenDs` \ [b,x] ->
 321
 322     -- build rest of the comprehesion
 323     dfListComp c_id b quals                     `thenDs` \ core_rest ->
 324
 325     -- build the pattern match
 326     matchSimply (Var x) (StmtCtxt ListComp)
 327                 pat core_rest (Var b)           `thenDs` \ core_expr ->
 328
 329     -- now build the outermost foldr, and return
 330     dsLookupGlobalId foldrName          `thenDs` \ foldr_id ->
 331     returnDs (
 332       Var foldr_id `App` Type x_ty
 333                    `App` Type b_ty
 334                    `App` mkLams [x, b] core_expr
 335                    `App` Var n_id
 336                    `App` core_list1
 337     )
 338 \end{code}
 339
 340 %************************************************************************
 341 %*                                                                      *
 342 \subsection[DsPArrComp]{Desugaring of array comprehensions}
 343 %*                                                                      *
 344 %************************************************************************
 345
 346 \begin{code}
 347
 348 -- entry point for desugaring a parallel array comprehension
 349 --
 350 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
 351 --
 352 dsPArrComp      :: [Stmt Id]
 353                 -> Type             -- Don't use; called with `undefined' below
 354                 -> DsM CoreExpr
 355 dsPArrComp qs _  =
 356   dsLookupGlobalId replicatePName                         `thenDs` \repP ->
 357   let unitArray = mkApps (Var repP) [Type unitTy,
 358                                      mkIntExpr 1,
 359                                      mkCoreTup []]
 360   in
 361   dePArrComp qs (noLoc (TuplePat [] Boxed)) unitArray
 362
 363 -- the work horse
 364 --
 365 dePArrComp :: [Stmt Id]
 366            -> LPat Id           -- the current generator pattern
 367            -> CoreExpr          -- the current generator expression
 368            -> DsM CoreExpr
 369 --
 370 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
 371 --
 372 dePArrComp [ResultStmt e'] pa cea =
 373   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 374   let ty = parrElemType cea
 375   in
 376   deLambda ty pa e'                                       `thenDs` \(clam,
 377                                                                      ty'e') ->
 378   returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
 379 --
 380 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
 381 --
 382 dePArrComp (ExprStmt b _ : qs) pa cea =
 383   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
 384   let ty = parrElemType cea
 385   in
 386   deLambda ty pa b                                        `thenDs` \(clam,_) ->
 387   dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
 388 --
 389 --  <<[:e' | p <- e, qs:]>> pa ea =
 390 --    let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
 391 --    in
 392 --    <<[:e' | qs:]>> (pa, p) (crossP ea ef)
 393 --
 394 dePArrComp (BindStmt p e : qs) pa cea =
 395   dsLookupGlobalId filterPName                    `thenDs` \filterP ->
 396   dsLookupGlobalId crossPName                     `thenDs` \crossP  ->
 397   dsLExpr e                                       `thenDs` \ce      ->
 398   let ty'cea = parrElemType cea
 399       ty'ce  = parrElemType ce
 400       false  = Var falseDataConId
 401       true   = Var trueDataConId
 402   in
 403   newSysLocalDs ty'ce                                     `thenDs` \v       ->
 404   matchSimply (Var v) (StmtCtxt PArrComp) p true false      `thenDs` \pred    ->
 405   let cef    = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
 406       ty'cef = ty'ce                            -- filterP preserves the type
 407       pa'    = noLoc (TuplePat [pa, p] Boxed)
 408   in
 409   dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
 410 --
 411 --  <<[:e' | let ds, qs:]>> pa ea =
 412 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
 413 --                    (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
 414 --  where
 415 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
 416 --
 417 dePArrComp (LetStmt ds : qs) pa cea =
 418   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 419   let xs     = map unLoc (collectGroupBinders ds)
 420       ty'cea = parrElemType cea
 421   in
 422   newSysLocalDs ty'cea                                    `thenDs` \v       ->
 423   dsLet ds (mkCoreTup (map Var xs))                       `thenDs` \clet    ->
 424   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
 425   let projBody = mkDsLet (NonRec let'v clet) $
 426                  mkCoreTup [Var v, Var let'v]
 427       errTy    = exprType projBody
 428       errMsg   = "DsListComp.dePArrComp: internal error!"
 429   in
 430   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 431   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr  `thenDs` \ccase   ->
 432   let pa'    = noLoc $ TuplePat [pa, noLoc (TuplePat (map nlVarPat xs) Boxed)] Boxed
 433       proj   = mkLams [v] ccase
 434   in
 435   dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
 436 --
 437 --  <<[:e' | qs | qss:]>> pa ea =
 438 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
 439 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
 440 --    where
 441 --      {x_1, ..., x_n} = DV (qs)
 442 --
 443 dePArrComp (ParStmt []             : qss2) pa cea = dePArrComp qss2 pa cea
 444 dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
 445   dsLookupGlobalId zipPName                               `thenDs` \zipP    ->
 446   let pa'     = noLoc $ TuplePat [pa, noLoc (TuplePat (map nlVarPat xs) Boxed)] Boxed
 447       ty'cea  = parrElemType cea
 448       resStmt = ResultStmt (noLoc $ ExplicitTuple (map nlHsVar xs) Boxed)
 449   in
 450   dsPArrComp (map unLoc qs ++ [resStmt]) undefined        `thenDs` \cqs     ->
 451   let ty'cqs = parrElemType cqs
 452       cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
 453   in
 454   dePArrComp (ParStmt qss : qss2) pa' cea'
 455
 456 -- generate Core corresponding to `\p -> e'
 457 --
 458 deLambda        :: Type                 -- type of the argument
 459                 -> LPat Id              -- argument pattern
 460                 -> LHsExpr Id           -- body
 461                 -> DsM (CoreExpr, Type)
 462 deLambda ty p e  =
 463   newSysLocalDs ty                                        `thenDs` \v       ->
 464   dsLExpr e                                               `thenDs` \ce      ->
 465   let errTy    = exprType ce
 466       errMsg   = "DsListComp.deLambda: internal error!"
 467   in
 468   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 469   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
 470   returnDs (mkLams [v] res, errTy)
 471
 472 -- obtain the element type of the parallel array produced by the given Core
 473 -- expression
 474 --
 475 parrElemType   :: CoreExpr -> Type
 476 parrElemType e  =
 477   case splitTyConApp_maybe (exprType e) of
 478     Just (tycon, [ty]) | tycon == parrTyCon -> ty
 479     _                                                     -> panic
 480       "DsListComp.parrElemType: not a parallel array type"
 481 \end{code}