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