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