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