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