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