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_IgnoreIfacePragmas, 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   |  opt_RulesOff || opt_IgnoreIfacePragmas     -- Either rules are switched off, or
  55                                                 --   we are ignoring what there are;
  56                                                 --   Either way foldr/build won't happen, so
  57                                                 --   use the more efficient Wadler-style desugaring
  58   || isParallelComp quals                       -- Foldr-style desugaring can't handle
  59                                                 --   parallel list comprehensions
  60   = deListComp quals (mkNilExpr elt_ty)
  61
  62   | otherwise           -- 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
 145 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
 146
 147 deListComp (ParStmt stmtss_w_bndrs : quals) list
 148   = mapDs do_list_comp stmtss_w_bndrs   `thenDs` \ exps ->
 149     mkZipBind qual_tys                  `thenDs` \ (zip_fn, zip_rhs) ->
 150
 151         -- Deal with [e | pat <- zip l1 .. ln] in example above
 152     deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
 153                    quals list
 154
 155   where
 156         bndrs_s = map snd stmtss_w_bndrs
 157
 158         -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
 159         pat      = TuplePat pats Boxed
 160         pats     = map mk_hs_tuple_pat bndrs_s
 161
 162         -- Types of (x1,..,xn), (y1,..,yn) etc
 163         qual_tys = map mk_bndrs_tys bndrs_s
 164
 165         do_list_comp (stmts, bndrs)
 166           = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
 167                        (mk_bndrs_tys bndrs)
 168
 169         mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
 170
 171         -- Last: the one to return
 172 deListComp [ResultStmt expr locn] list  -- Figure 7.4, SLPJ, p 135, rule C above
 173   = dsExpr expr                 `thenDs` \ core_expr ->
 174     returnDs (mkConsExpr (exprType core_expr) core_expr list)
 175
 176         -- Non-last: must be a guard
 177 deListComp (ExprStmt guard ty locn : quals) list        -- rule B above
 178   = dsExpr guard                `thenDs` \ core_guard ->
 179     deListComp quals list       `thenDs` \ core_rest ->
 180     returnDs (mkIfThenElse core_guard core_rest list)
 181
 182 -- [e | let B, qs] = let B in [e | qs]
 183 deListComp (LetStmt binds : quals) list
 184   = deListComp quals list       `thenDs` \ core_rest ->
 185     dsLet binds core_rest
 186
 187 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
 188   = dsExpr list1                    `thenDs` \ core_list1 ->
 189     deBindComp pat core_list1 quals core_list2
 190 \end{code}
 191
 192
 193 \begin{code}
 194 deBindComp pat core_list1 quals core_list2
 195   = let
 196         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
 197
 198         -- u1_ty is a [alpha] type, and u2_ty = alpha
 199         u2_ty = hsPatType pat
 200
 201         res_ty = exprType core_list2
 202         h_ty   = u1_ty `mkFunTy` res_ty
 203     in
 204     newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]  `thenDs` \ [h, u1, u2, u3] ->
 205
 206     -- the "fail" value ...
 207     let
 208         core_fail   = App (Var h) (Var u3)
 209         letrec_body = App (Var h) core_list1
 210     in
 211     deListComp quals core_fail                  `thenDs` \ rest_expr ->
 212     matchSimply (Var u2) (StmtCtxt ListComp) pat
 213                 rest_expr core_fail             `thenDs` \ core_match ->
 214     let
 215         rhs = Lam u1 $
 216               Case (Var u1) u1 [(DataAlt nilDataCon,  [],       core_list2),
 217                                 (DataAlt consDataCon, [u2, u3], core_match)]
 218     in
 219     returnDs (Let (Rec [(h, rhs)]) letrec_body)
 220 \end{code}
 221
 222
 223 \begin{code}
 224 mkZipBind :: [Type] -> DsM (Id, CoreExpr)
 225 -- mkZipBind [t1, t2]
 226 -- = (zip, \as1:[t1] as2:[t2]
 227 --         -> case as1 of
 228 --              [] -> []
 229 --              (a1:as'1) -> case as2 of
 230 --                              [] -> []
 231 --                              (a2:as'2) -> (a2,a2) : zip as'1 as'2)]
 232
 233 mkZipBind elt_tys
 234   = mapDs newSysLocalDs  list_tys       `thenDs` \ ass ->
 235     mapDs newSysLocalDs  elt_tys        `thenDs` \ as' ->
 236     mapDs newSysLocalDs  list_tys       `thenDs` \ as's ->
 237     newSysLocalDs zip_fn_ty             `thenDs` \ zip_fn ->
 238     let
 239         inner_rhs = mkConsExpr ret_elt_ty
 240                         (mkCoreTup (map Var as'))
 241                         (mkVarApps (Var zip_fn) as's)
 242         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
 243     in
 244     returnDs (zip_fn, mkLams ass zip_body)
 245   where
 246     list_tys   = map mkListTy elt_tys
 247     ret_elt_ty = mkCoreTupTy elt_tys
 248     zip_fn_ty  = mkFunTys list_tys (mkListTy ret_elt_ty)
 249
 250     mk_case (as, a', as') rest
 251           = Case (Var as) as [(DataAlt nilDataCon,  [],        mkNilExpr ret_elt_ty),
 252                               (DataAlt consDataCon, [a', as'], rest)]
 253
 254 -- Helper functions that makes an HsTuple only for non-1-sized tuples
 255 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
 256 mk_hs_tuple_expr []   = HsVar unitDataConId
 257 mk_hs_tuple_expr [id] = HsVar id
 258 mk_hs_tuple_expr ids  = ExplicitTuple [ HsVar i | i <- ids ] Boxed
 259
 260 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
 261 mk_hs_tuple_pat [b] = VarPat b
 262 mk_hs_tuple_pat bs  = TuplePat (map VarPat bs) Boxed
 263 \end{code}
 264
 265
 266 %************************************************************************
 267 %*                                                                      *
 268 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 269 %*                                                                      *
 270 %************************************************************************
 271
 272 @dfListComp@ are the rules used with foldr/build turned on:
 273
 274 \begin{verbatim}
 275 TE[ e | ]            c n = c e n
 276 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
 277 TE[ e | p <- l , q ] c n = let
 278                                 f = \ x b -> case x of
 279                                                   p -> TE[ e | q ] c b
 280                                                   _ -> b
 281                            in
 282                            foldr f n l
 283 \end{verbatim}
 284
 285 \begin{code}
 286 dfListComp :: Id -> Id                  -- 'c' and 'n'
 287            -> [TypecheckedStmt]         -- the rest of the qual's
 288            -> DsM CoreExpr
 289
 290         -- Last: the one to return
 291 dfListComp c_id n_id [ResultStmt expr locn]
 292   = dsExpr expr                 `thenDs` \ core_expr ->
 293     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 294
 295         -- Non-last: must be a guard
 296 dfListComp c_id n_id (ExprStmt guard ty locn  : quals)
 297   = dsExpr guard                                `thenDs` \ core_guard ->
 298     dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 299     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 300
 301 dfListComp c_id n_id (LetStmt binds : quals)
 302   -- new in 1.3, local bindings
 303   = dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 304     dsLet binds core_rest
 305
 306 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
 307     -- evaluate the two lists
 308   = dsExpr list1                                `thenDs` \ core_list1 ->
 309
 310     -- find the required type
 311     let x_ty   = hsPatType pat
 312         b_ty   = idType n_id
 313     in
 314
 315     -- create some new local id's
 316     newSysLocalsDs [b_ty,x_ty]                  `thenDs` \ [b,x] ->
 317
 318     -- build rest of the comprehesion
 319     dfListComp c_id b quals                     `thenDs` \ core_rest ->
 320
 321     -- build the pattern match
 322     matchSimply (Var x) (StmtCtxt ListComp)
 323                 pat core_rest (Var b)           `thenDs` \ core_expr ->
 324
 325     -- now build the outermost foldr, and return
 326     dsLookupGlobalId foldrName          `thenDs` \ foldr_id ->
 327     returnDs (
 328       Var foldr_id `App` Type x_ty
 329                    `App` Type b_ty
 330                    `App` mkLams [x, b] core_expr
 331                    `App` Var n_id
 332                    `App` core_list1
 333     )
 334 \end{code}
 335
 336 %************************************************************************
 337 %*                                                                      *
 338 \subsection[DsPArrComp]{Desugaring of array comprehensions}
 339 %*                                                                      *
 340 %************************************************************************
 341
 342 \begin{code}
 343
 344 -- entry point for desugaring a parallel array comprehension
 345 --
 346 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
 347 --
 348 dsPArrComp      :: [TypecheckedStmt]
 349                 -> Type             -- Don't use; called with `undefined' below
 350                 -> DsM CoreExpr
 351 dsPArrComp qs _  =
 352   dsLookupGlobalId replicatePName                         `thenDs` \repP ->
 353   let unitArray = mkApps (Var repP) [Type unitTy,
 354                                      mkIntExpr 1,
 355                                      mkCoreTup []]
 356   in
 357   dePArrComp qs (TuplePat [] Boxed) unitArray
 358
 359 -- the work horse
 360 --
 361 dePArrComp :: [TypecheckedStmt]
 362            -> TypecheckedPat            -- the current generator pattern
 363            -> CoreExpr                  -- the current generator expression
 364            -> DsM CoreExpr
 365 --
 366 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
 367 --
 368 dePArrComp [ResultStmt e' _] pa cea =
 369   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 370   let ty = parrElemType cea
 371   in
 372   deLambda ty pa e'                                       `thenDs` \(clam,
 373                                                                      ty'e') ->
 374   returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
 375 --
 376 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
 377 --
 378 dePArrComp (ExprStmt b _ _ : qs) pa cea =
 379   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
 380   let ty = parrElemType cea
 381   in
 382   deLambda ty pa b                                        `thenDs` \(clam,_) ->
 383   dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
 384 --
 385 --  <<[:e' | p <- e, qs:]>> pa ea =
 386 --    let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
 387 --    in
 388 --    <<[:e' | qs:]>> (pa, p) (crossP ea ef)
 389 --
 390 dePArrComp (BindStmt p e _ : qs) pa cea =
 391   dsLookupGlobalId filterPName                    `thenDs` \filterP ->
 392   dsLookupGlobalId crossPName                     `thenDs` \crossP  ->
 393   dsExpr e                                        `thenDs` \ce      ->
 394   let ty'cea = parrElemType cea
 395       ty'ce  = parrElemType ce
 396       false  = Var falseDataConId
 397       true   = Var trueDataConId
 398   in
 399   newSysLocalDs ty'ce                                     `thenDs` \v       ->
 400   matchSimply (Var v) (StmtCtxt PArrComp) p true false      `thenDs` \pred    ->
 401   let cef    = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
 402       ty'cef = ty'ce                            -- filterP preserves the type
 403       pa'    = TuplePat [pa, p] Boxed
 404   in
 405   dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
 406 --
 407 --  <<[:e' | let ds, qs:]>> pa ea =
 408 --    <<[:e' | qs:]>> (pa, (x_1, ..., x_n))
 409 --                    (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
 410 --  where
 411 --    {x_1, ..., x_n} = DV (ds)         -- Defined Variables
 412 --
 413 dePArrComp (LetStmt ds : qs) pa cea =
 414   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 415   let xs     = collectHsBinders ds
 416       ty'cea = parrElemType cea
 417   in
 418   newSysLocalDs ty'cea                                    `thenDs` \v       ->
 419   dsLet ds (mkCoreTup (map Var xs))                       `thenDs` \clet    ->
 420   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
 421   let projBody = mkDsLet (NonRec let'v clet) $
 422                  mkCoreTup [Var v, Var let'v]
 423       errTy    = exprType projBody
 424       errMsg   = "DsListComp.dePArrComp: internal error!"
 425   in
 426   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 427   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr  `thenDs` \ccase   ->
 428   let pa'    = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
 429       proj   = mkLams [v] ccase
 430   in
 431   dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
 432 --
 433 --  <<[:e' | qs | qss:]>> pa ea =
 434 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
 435 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
 436 --    where
 437 --      {x_1, ..., x_n} = DV (qs)
 438 --
 439 dePArrComp (ParStmt []             : qss2) pa cea = dePArrComp qss2 pa cea
 440 dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
 441   dsLookupGlobalId zipPName                               `thenDs` \zipP    ->
 442   let pa'     = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
 443       ty'cea  = parrElemType cea
 444       resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
 445   in
 446   dsPArrComp (qs ++ [resStmt]) undefined                  `thenDs` \cqs     ->
 447   let ty'cqs = parrElemType cqs
 448       cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
 449   in
 450   dePArrComp (ParStmt qss : qss2) pa' cea'
 451
 452 -- generate Core corresponding to `\p -> e'
 453 --
 454 deLambda        :: Type                 -- type of the argument
 455                 -> TypecheckedPat       -- argument pattern
 456                 -> TypecheckedHsExpr    -- body
 457                 -> DsM (CoreExpr, Type)
 458 deLambda ty p e  =
 459   newSysLocalDs ty                                        `thenDs` \v       ->
 460   dsExpr e                                                `thenDs` \ce      ->
 461   let errTy    = exprType ce
 462       errMsg   = "DsListComp.deLambda: internal error!"
 463   in
 464   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 465   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
 466   returnDs (mkLams [v] res, errTy)
 467
 468 -- obtain the element type of the parallel array produced by the given Core
 469 -- expression
 470 --
 471 parrElemType   :: CoreExpr -> Type
 472 parrElemType e  =
 473   case splitTyConApp_maybe (exprType e) of
 474     Just (tycon, [ty]) | tyConName tycon == parrTyConName -> ty
 475     _                                                     -> panic
 476       "DsListComp.parrElemType: not a parallel array type"
 477 \end{code}