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