ghc/compiler/deSugar/DsListComp.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[DsListComp]{Desugaring list comprehensions}
   5
   6 \begin{code}
   7 module DsListComp ( dsListComp ) where
   8
   9 #include "HsVersions.h"
  10
  11 import {-# SOURCE #-} DsExpr ( dsExpr, dsLet )
  12
  13 import BasicTypes       ( Boxity(..) )
  14 import HsSyn            ( OutPat(..), HsExpr(..), Stmt(..), HsMatchContext(..) )
  15 import TcHsSyn          ( TypecheckedStmt )
  16 import DsHsSyn          ( outPatType )
  17 import CoreSyn
  18
  19 import DsMonad          -- the monadery used in the desugarer
  20 import DsUtils
  21
  22 import CmdLineOpts      ( opt_FoldrBuildOn )
  23 import CoreUtils        ( exprType, mkIfThenElse )
  24 import Id               ( idType )
  25 import Var              ( Id )
  26 import Type             ( mkTyVarTy, mkFunTys, mkFunTy, Type )
  27 import TysPrim          ( alphaTyVar )
  28 import TysWiredIn       ( nilDataCon, consDataCon, unitDataConId, tupleCon, mkListTy, mkTupleTy )
  29 import Match            ( matchSimply )
  30 import PrelNames        ( foldrName, buildName )
  31 import SrcLoc           ( noSrcLoc )
  32 import List             ( zip4 )
  33 \end{code}
  34
  35 List comprehensions may be desugared in one of two ways: ``ordinary''
  36 (as you would expect if you read SLPJ's book) and ``with foldr/build
  37 turned on'' (if you read Gill {\em et al.}'s paper on the subject).
  38
  39 There will be at least one ``qualifier'' in the input.
  40
  41 \begin{code}
  42 dsListComp :: [TypecheckedStmt]
  43            -> Type              -- Type of list elements
  44            -> DsM CoreExpr
  45
  46 dsListComp quals elt_ty
  47   |  not opt_FoldrBuildOn                -- Be boring
  48   || isParallelComp quals
  49   = deListComp quals (mkNilExpr elt_ty)
  50
  51   | otherwise                            -- foldr/build lives!
  52   = newTyVarsDs [alphaTyVar]    `thenDs` \ [n_tyvar] ->
  53     let
  54         n_ty = mkTyVarTy n_tyvar
  55         c_ty = mkFunTys [elt_ty, n_ty] n_ty
  56     in
  57     newSysLocalsDs [c_ty,n_ty]          `thenDs` \ [c, n] ->
  58     dfListComp c n quals                `thenDs` \ result ->
  59     dsLookupGlobalValue buildName       `thenDs` \ build_id ->
  60     returnDs (Var build_id `App` Type elt_ty
  61                            `App` mkLams [n_tyvar, c, n] result)
  62
  63   where isParallelComp (ParStmtOut bndrstmtss : _) = True
  64         isParallelComp _                           = False
  65 \end{code}
  66
  67 %************************************************************************
  68 %*                                                                      *
  69 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
  70 %*                                                                      *
  71 %************************************************************************
  72
  73 Just as in Phil's chapter~7 in SLPJ, using the rules for
  74 optimally-compiled list comprehensions.  This is what Kevin followed
  75 as well, and I quite happily do the same.  The TQ translation scheme
  76 transforms a list of qualifiers (either boolean expressions or
  77 generators) into a single expression which implements the list
  78 comprehension.  Because we are generating 2nd-order polymorphic
  79 lambda-calculus, calls to NIL and CONS must be applied to a type
  80 argument, as well as their usual value arguments.
  81 \begin{verbatim}
  82 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
  83
  84 (Rule C)
  85 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
  86
  87 (Rule B)
  88 TQ << [ e | b , qs ] ++ L >> =
  89     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
  90
  91 (Rule A')
  92 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
  93   letrec
  94     h = \ u1 ->
  95           case u1 of
  96             []        ->  TE << L2 >>
  97             (u2 : u3) ->
  98                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
  99                     [] (h u3)
 100   in
 101     h ( TE << L1 >> )
 102
 103 "h", "u1", "u2", and "u3" are new variables.
 104 \end{verbatim}
 105
 106 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
 107 is the TE translation scheme.  Note that we carry around the @L@ list
 108 already desugared.  @dsListComp@ does the top TE rule mentioned above.
 109
 110 To the above, we add an additional rule to deal with parallel list
 111 comprehensions.  The translation goes roughly as follows:
 112      [ e | p1 <- e11, let v1 = e12, p2 <- e13
 113          | q1 <- e21, let v2 = e22, q2 <- e23]
 114      =>
 115      [ e | ((p1,v1,p2), (q1,v2,q2)) <-
 116                zip [(p1,v1,p2) | p1 <- e11, let v1 = e12, p2 <- e13]
 117                    [(q1,v2,q2) | q1 <- e21, let v2 = e22, q2 <- e23]]
 118 In the translation below, the ParStmtOut branch translates each parallel branch
 119 into a sub-comprehension, and desugars each independently.  The resulting lists
 120 are fed to a zip function, we create a binding for all the variables bound in all
 121 the comprehensions, and then we hand things off the the desugarer for bindings.
 122 The zip function is generated here a) because it's small, and b) because then we
 123 don't have to deal with arbitrary limits on the number of zip functions in the
 124 prelude, nor which library the zip function came from.
 125 The introduced tuples are Boxed, but only because I couldn't get it to work
 126 with the Unboxed variety.
 127
 128 \begin{code}
 129
 130 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
 131
 132 deListComp (ParStmtOut bndrstmtss : quals) list
 133   = mapDs doListComp qualss     `thenDs` \ exps ->
 134     mapDs genAS  bndrss         `thenDs` \ ass ->
 135     mapDs genA   bndrss         `thenDs` \ as ->
 136     mapDs genAS' bndrss         `thenDs` \ as's ->
 137     let retTy = myTupleTy Boxed (length bndrss) qualTys
 138         zipTy = foldr mkFunTy (mkListTy retTy) (map mkListTy qualTys)
 139     in
 140     newSysLocalDs zipTy         `thenDs` \ zipFn ->
 141     let target = mkConsExpr retTy (mkTupleExpr as) (foldl App (Var zipFn) (map Var as's))
 142         zipExp = mkLet zipFn (zip4 bndrss ass as as's) exps target
 143     in
 144     deBindComp pat zipExp quals list
 145   where (bndrss, stmtss) = unzip bndrstmtss
 146         pats = map (\ps -> mkTuplePat (map VarPat ps)) bndrss
 147         mkTuplePat [p] = p
 148         mkTuplePat ps  = TuplePat ps Boxed
 149         pat  = TuplePat pats Boxed
 150
 151         qualss = map mkQuals bndrstmtss
 152         mkQuals (bndrs, stmts) = (bndrs, stmts ++ [ExprStmt (myTupleExpr bndrs) noSrcLoc])
 153
 154         qualTys = map mkBndrsTy bndrss
 155         mkBndrsTy bndrs = myTupleTy Boxed (length bndrs) (map idType bndrs)
 156
 157         doListComp (bndrs, stmts)
 158           = dsListComp stmts (mkBndrsTy bndrs)
 159         genA   bndrs = newSysLocalDs (mkBndrsTy bndrs)
 160         genAS  bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
 161         genAS' bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
 162
 163         mkLet zipFn vars exps target
 164           = Let (Rec [(zipFn,
 165                        foldr Lam (mkBody target vars) (map getAs vars))])
 166                 (foldl App (Var zipFn) exps)
 167         getAs (_, as, _, _) = as
 168         mkBody target vars
 169           = foldr mkCase (foldr mkTuplCase target vars) vars
 170         mkCase (ps, as, a, as') rest
 171           = Case (Var as) as [(DataAlt nilDataCon, [], mkConApp nilDataCon []),
 172                               (DataAlt consDataCon, [a, as'], rest)]
 173         mkTuplCase ([p], as, a, as') rest
 174           = App (Lam p rest) (Var a)
 175         mkTuplCase (ps, as, a, as') rest
 176           = Case (Var a) a [(DataAlt (tupleCon Boxed (length ps)), ps, rest)]
 177
 178         myTupleTy boxity arity [ty] = ty
 179         myTupleTy boxity arity tys  = mkTupleTy boxity arity tys
 180         myTupleExpr []   = HsVar unitDataConId
 181         myTupleExpr [id] = HsVar id
 182         myTupleExpr ids  = ExplicitTuple [ HsVar i | i <- ids ] Boxed
 183
 184         -- Last: the one to return
 185 deListComp [ExprStmt expr locn] list    -- Figure 7.4, SLPJ, p 135, rule C above
 186   = dsExpr expr                 `thenDs` \ core_expr ->
 187     returnDs (mkConsExpr (exprType core_expr) core_expr list)
 188
 189         -- Non-last: must be a guard
 190 deListComp (ExprStmt guard locn : quals) list   -- rule B above
 191   = dsExpr guard                `thenDs` \ core_guard ->
 192     deListComp quals list       `thenDs` \ core_rest ->
 193     returnDs (mkIfThenElse core_guard core_rest list)
 194
 195 -- [e | let B, qs] = let B in [e | qs]
 196 deListComp (LetStmt binds : quals) list
 197   = deListComp quals list       `thenDs` \ core_rest ->
 198     dsLet binds core_rest
 199
 200 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
 201   = dsExpr list1                    `thenDs` \ core_list1 ->
 202     deBindComp pat core_list1 quals core_list2
 203
 204 deBindComp pat core_list1 quals core_list2
 205   = let
 206         u3_ty@u1_ty = exprType core_list1       -- two names, same thing
 207
 208         -- u1_ty is a [alpha] type, and u2_ty = alpha
 209         u2_ty = outPatType pat
 210
 211         res_ty = exprType core_list2
 212         h_ty   = u1_ty `mkFunTy` res_ty
 213     in
 214     newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]  `thenDs` \ [h, u1, u2, u3] ->
 215
 216     -- the "fail" value ...
 217     let
 218         core_fail   = App (Var h) (Var u3)
 219         letrec_body = App (Var h) core_list1
 220     in
 221     deListComp quals core_fail                  `thenDs` \ rest_expr ->
 222     matchSimply (Var u2) ListComp pat
 223                 rest_expr core_fail             `thenDs` \ core_match ->
 224     let
 225         rhs = Lam u1 $
 226               Case (Var u1) u1 [(DataAlt nilDataCon,  [],       core_list2),
 227                                 (DataAlt consDataCon, [u2, u3], core_match)]
 228     in
 229     returnDs (Let (Rec [(h, rhs)]) letrec_body)
 230 \end{code}
 231
 232
 233 %************************************************************************
 234 %*                                                                      *
 235 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 236 %*                                                                      *
 237 %************************************************************************
 238
 239 @dfListComp@ are the rules used with foldr/build turned on:
 240
 241 \begin{verbatim}
 242 TE[ e | ]            c n = c e n
 243 TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
 244 TE[ e | p <- l , q ] c n = let
 245                                 f = \ x b -> case x of
 246                                                   p -> TE[ e | q ] c b
 247                                                   _ -> b
 248                            in
 249                            foldr f n l
 250 \end{verbatim}
 251
 252 \begin{code}
 253 dfListComp :: Id -> Id                  -- 'c' and 'n'
 254            -> [TypecheckedStmt]         -- the rest of the qual's
 255            -> DsM CoreExpr
 256
 257         -- Last: the one to return
 258 dfListComp c_id n_id [ExprStmt expr locn]
 259   = dsExpr expr                 `thenDs` \ core_expr ->
 260     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 261
 262         -- Non-last: must be a guard
 263 dfListComp c_id n_id (ExprStmt guard locn  : quals)
 264   = dsExpr guard                                `thenDs` \ core_guard ->
 265     dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 266     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 267
 268 dfListComp c_id n_id (LetStmt binds : quals)
 269   -- new in 1.3, local bindings
 270   = dfListComp c_id n_id quals  `thenDs` \ core_rest ->
 271     dsLet binds core_rest
 272
 273 dfListComp c_id n_id (BindStmt pat list1 locn : quals)
 274     -- evaluate the two lists
 275   = dsExpr list1                                `thenDs` \ core_list1 ->
 276
 277     -- find the required type
 278     let x_ty   = outPatType pat
 279         b_ty   = idType n_id
 280     in
 281
 282     -- create some new local id's
 283     newSysLocalsDs [b_ty,x_ty]                  `thenDs` \ [b,x] ->
 284
 285     -- build rest of the comprehesion
 286     dfListComp c_id b quals                     `thenDs` \ core_rest ->
 287
 288     -- build the pattern match
 289     matchSimply (Var x) ListComp pat core_rest (Var b)  `thenDs` \ core_expr ->
 290
 291     -- now build the outermost foldr, and return
 292     dsLookupGlobalValue foldrName               `thenDs` \ foldr_id ->
 293     returnDs (
 294       Var foldr_id `App` Type x_ty
 295                    `App` Type b_ty
 296                    `App` mkLams [x, b] core_expr
 297                    `App` Var n_id
 298                    `App` core_list1
 299     )
 300 \end{code}
 301
 302