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