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