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