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