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