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