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,
  34                           mkListTy, mkTupleTy )
  35 import Match            ( matchSimply )
  36 import PrelNames        ( foldrName, buildName, replicatePName, mapPName,
  37                           filterPName, 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
 233                         (mkCoreTup (map Var as'))
 234                         (mkVarApps (Var zip_fn) as's)
 235         zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
 236     in
 237     returnDs (zip_fn, mkLams ass zip_body)
 238   where
 239     list_tys   = map mkListTy elt_tys
 240     ret_elt_ty = mk_tuple_ty elt_tys
 241     zip_fn_ty  = mkFunTys list_tys (mkListTy ret_elt_ty)
 242
 243     mk_case (as, a', as') rest
 244           = Case (Var as) as [(DataAlt nilDataCon,  [],        mkNilExpr ret_elt_ty),
 245                               (DataAlt consDataCon, [a', as'], rest)]
 246
 247 -- Helper function
 248 mk_tuple_ty :: [Type] -> Type
 249 mk_tuple_ty [ty] = ty
 250 mk_tuple_ty tys  = mkTupleTy Boxed (length tys) tys
 251
 252 -- Helper functions that makes an HsTuple only for non-1-sized tuples
 253 mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
 254 mk_hs_tuple_expr []   = HsVar unitDataConId
 255 mk_hs_tuple_expr [id] = HsVar id
 256 mk_hs_tuple_expr ids  = ExplicitTuple [ HsVar i | i <- ids ] Boxed
 257
 258 mk_hs_tuple_pat :: [Id] -> TypecheckedPat
 259 mk_hs_tuple_pat [b] = VarPat b
 260 mk_hs_tuple_pat bs  = TuplePat (map VarPat bs) Boxed
 261 \end{code}
 262
 263
 264 %************************************************************************
 265 %*                                                                      *
 266 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 267 %*                                                                      *
 268 %************************************************************************
 269
 270 @dfListComp@ are the rules used with foldr/build turned on:
 271
 272 \begin{verbatim}
 273 TE[ e | ]            c n = c e n
 274 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
 275 TE[ e | p <- l , q ] c n = let
 276                                 f = \ x b -> case x of
 277                                                   p -> TE[ e | q ] c b
 278                                                   _ -> b
 279                            in
 280                            foldr f n l
 281 \end{verbatim}
 282
 283 \begin{code}
 284 dfListComp :: Id -> Id                  -- 'c' and 'n'
 285            -> [TypecheckedStmt]         -- the rest of the qual's
 286            -> DsM CoreExpr
 287
 288         -- Last: the one to return
 289 dfListComp c_id n_id [ResultStmt expr locn]
 290   = dsExpr expr                 `thenDs` \ core_expr ->
 291     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 292
 293         -- Non-last: must be a guard
 294 dfListComp c_id n_id (ExprStmt guard ty locn  : quals)
 295   = dsExpr guard                                `thenDs` \ core_guard ->
 296     dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 297     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 298
 299 dfListComp c_id n_id (LetStmt binds : quals)
 300   -- new in 1.3, local bindings
 301   = dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 302     dsLet binds core_rest
 303
 304 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
 305     -- evaluate the two lists
 306   = dsExpr list1                                `thenDs` \ core_list1 ->
 307
 308     -- find the required type
 309     let x_ty   = hsPatType pat
 310         b_ty   = idType n_id
 311     in
 312
 313     -- create some new local id's
 314     newSysLocalsDs [b_ty,x_ty]                  `thenDs` \ [b,x] ->
 315
 316     -- build rest of the comprehesion
 317     dfListComp c_id b quals                     `thenDs` \ core_rest ->
 318
 319     -- build the pattern match
 320     matchSimply (Var x) (StmtCtxt ListComp)
 321                 pat core_rest (Var b)           `thenDs` \ core_expr ->
 322
 323     -- now build the outermost foldr, and return
 324     dsLookupGlobalId foldrName          `thenDs` \ foldr_id ->
 325     returnDs (
 326       Var foldr_id `App` Type x_ty
 327                    `App` Type b_ty
 328                    `App` mkLams [x, b] core_expr
 329                    `App` Var n_id
 330                    `App` core_list1
 331     )
 332 \end{code}
 333
 334 %************************************************************************
 335 %*                                                                      *
 336 \subsection[DsPArrComp]{Desugaring of array comprehensions}
 337 %*                                                                      *
 338 %************************************************************************
 339
 340 \begin{code}
 341
 342 -- entry point for desugaring a parallel array comprehension
 343 --
 344 --   [:e | qss:] = <<[:e | qss:]>> () [:():]
 345 --
 346 dsPArrComp      :: [TypecheckedStmt]
 347                 -> Type             -- Don't use; called with `undefined' below
 348                 -> DsM CoreExpr
 349 dsPArrComp qs _  =
 350   dsLookupGlobalId replicatePName                         `thenDs` \repP ->
 351   let unitArray = mkApps (Var repP) [Type unitTy,
 352                                      mkIntExpr 1,
 353                                      mkCoreTup []]
 354   in
 355   dePArrComp qs (TuplePat [] Boxed) unitArray
 356
 357 -- the work horse
 358 --
 359 dePArrComp :: [TypecheckedStmt]
 360            -> TypecheckedPat            -- the current generator pattern
 361            -> CoreExpr                  -- the current generator expression
 362            -> DsM CoreExpr
 363 --
 364 --  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
 365 --
 366 dePArrComp [ResultStmt e' _] pa cea =
 367   dsLookupGlobalId mapPName                               `thenDs` \mapP    ->
 368   let ty = parrElemType cea
 369   in
 370   deLambda ty pa e'                                       `thenDs` \(clam,
 371                                                                      ty'e') ->
 372   returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
 373 --
 374 --  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
 375 --
 376 dePArrComp (ExprStmt b _ _ : qs) pa cea =
 377   dsLookupGlobalId filterPName                    `thenDs` \filterP  ->
 378   let ty = parrElemType cea
 379   in
 380   deLambda ty pa b                                        `thenDs` \(clam,_) ->
 381   dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
 382 --
 383 --  <<[:e' | p <- e, qs:]>> pa ea =
 384 --    let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
 385 --    in
 386 --    <<[:e' | qs:]>> (pa, p) (crossP ea ef)
 387 --
 388 dePArrComp (BindStmt p e _ : qs) pa cea =
 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 falseDataConId
 395       true   = Var trueDataConId
 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 (mkCoreTup (map Var xs))                       `thenDs` \clet    ->
 418   newSysLocalDs (exprType clet)                           `thenDs` \let'v   ->
 419   let projBody = mkDsLet (NonRec let'v clet) $
 420                  mkCoreTup [Var v, Var let'v]
 421       errTy    = exprType projBody
 422       errMsg   = "DsListComp.dePArrComp: internal error!"
 423   in
 424   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 425   matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr  `thenDs` \ccase   ->
 426   let pa'    = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
 427       proj   = mkLams [v] ccase
 428   in
 429   dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
 430 --
 431 --  <<[:e' | qs | qss:]>> pa ea =
 432 --    <<[:e' | qss:]>> (pa, (x_1, ..., x_n))
 433 --                     (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
 434 --    where
 435 --      {x_1, ..., x_n} = DV (qs)
 436 --
 437 dePArrComp (ParStmtOut []             : qss2) pa cea = dePArrComp qss2 pa cea
 438 dePArrComp (ParStmtOut ((xs, qs):qss) : qss2) pa cea =
 439   dsLookupGlobalId zipPName                               `thenDs` \zipP    ->
 440   let pa'     = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
 441       ty'cea  = parrElemType cea
 442       resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
 443   in
 444   dsPArrComp (qs ++ [resStmt]) undefined                  `thenDs` \cqs     ->
 445   let ty'cqs = parrElemType cqs
 446       cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
 447   in
 448   dePArrComp (ParStmtOut qss : qss2) pa' cea'
 449
 450 -- generate Core corresponding to `\p -> e'
 451 --
 452 deLambda        :: Type                 -- type of the argument
 453                 -> TypecheckedPat       -- argument pattern
 454                 -> TypecheckedHsExpr    -- body
 455                 -> DsM (CoreExpr, Type)
 456 deLambda ty p e  =
 457   newSysLocalDs ty                                        `thenDs` \v       ->
 458   dsExpr e                                                `thenDs` \ce      ->
 459   let errTy    = exprType ce
 460       errMsg   = "DsListComp.deLambda: internal error!"
 461   in
 462   mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
 463   matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr       `thenDs` \res     ->
 464   returnDs (mkLams [v] res, errTy)
 465
 466 -- obtain the element type of the parallel array produced by the given Core
 467 -- expression
 468 --
 469 parrElemType   :: CoreExpr -> Type
 470 parrElemType e  =
 471   case splitTyConApp_maybe (exprType e) of
 472     Just (tycon, [ty]) | tyConName tycon == parrTyConName -> ty
 473     _                                                     -> panic
 474       "DsListComp.parrElemType: not a parallel array type"
 475 \end{code}