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 HsSyn            ( Stmt(..), HsExpr )
  14 import TcHsSyn          ( TypecheckedStmt, TypecheckedHsExpr )
  15 import DsHsSyn          ( 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        ( coreExprType )
  23 import Var              ( Id, TyVar )
  24 import Const            ( Con(..) )
  25 import PrelInfo         ( foldrId )
  26 import Type             ( mkTyVarTy, mkForAllTy, mkFunTys, mkFunTy, Type )
  27 import TysPrim          ( alphaTyVar, alphaTy )
  28 import TysWiredIn       ( nilDataCon, consDataCon, listTyCon )
  29 import Match            ( matchSimply )
  30 import Outputable
  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   = deListComp quals nil_expr
  47
  48   | otherwise                            -- foldr/build lives!
  49   = newTyVarsDs [alphaTyVar]    `thenDs` \ [n_tyvar] ->
  50     let
  51         alpha_to_alpha = alphaTy `mkFunTy` alphaTy
  52
  53         n_ty = mkTyVarTy n_tyvar
  54         c_ty = mkFunTys [elt_ty, n_ty] n_ty
  55         g_ty = mkForAllTy alphaTyVar (
  56                 (elt_ty `mkFunTy` alpha_to_alpha)
  57                 `mkFunTy`
  58                 alpha_to_alpha
  59            )
  60     in
  61     newSysLocalsDs [c_ty,n_ty,g_ty]  `thenDs` \ [c, n, g] ->
  62
  63     dfListComp  c_ty c
  64                 n_ty n
  65                 quals       `thenDs` \ result ->
  66
  67     returnDs (mkBuild elt_ty n_tyvar c n g result)
  68   where
  69     nil_expr = mkNilExpr elt_ty
  70 \end{code}
  71
  72 %************************************************************************
  73 %*                                                                      *
  74 \subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
  75 %*                                                                      *
  76 %************************************************************************
  77
  78 Just as in Phil's chapter~7 in SLPJ, using the rules for
  79 optimally-compiled list comprehensions.  This is what Kevin followed
  80 as well, and I quite happily do the same.  The TQ translation scheme
  81 transforms a list of qualifiers (either boolean expressions or
  82 generators) into a single expression which implements the list
  83 comprehension.  Because we are generating 2nd-order polymorphic
  84 lambda-calculus, calls to NIL and CONS must be applied to a type
  85 argument, as well as their usual value arguments.
  86 \begin{verbatim}
  87 TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>
  88
  89 (Rule C)
  90 TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>
  91
  92 (Rule B)
  93 TQ << [ e | b , qs ] ++ L >> =
  94     if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>
  95
  96 (Rule A')
  97 TQ << [ e | p <- L1, qs ]  ++  L2 >> =
  98   letrec
  99     h = \ u1 ->
 100           case u1 of
 101             []        ->  TE << L2 >>
 102             (u2 : u3) ->
 103                   (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
 104                     [] (h u3)
 105   in
 106     h ( TE << L1 >> )
 107
 108 "h", "u1", "u2", and "u3" are new variables.
 109 \end{verbatim}
 110
 111 @deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
 112 is the TE translation scheme.  Note that we carry around the @L@ list
 113 already desugared.  @dsListComp@ does the top TE rule mentioned above.
 114
 115 deListComp :: [TypecheckedStmt]
 116            -> CoreExpr -> CoreExpr      -- Cons and nil resp; can be copied freely
 117            -> DsM CoreExpr
 118
 119 deListComp [ReturnStmt expr] cons nil
 120   = dsExpr expr                 `thenDs` \ expr' ->
 121     returnDs (mkApps cons [expr', nil])
 122
 123 deListComp (GuardStmt guard locn : quals) cons nil
 124   = dsExpr guard                `thenDs` \ guard' ->
 125     deListComp quals cons nil   `thenDs` \ rest' ->
 126     returnDs (mkIfThenElse guard' rest' nil)
 127
 128 deListComp (LetStmt binds : quals) cons nil
 129   = deListComp quals cons nil           `thenDs` \ rest' ->
 130     dsLet binds rest'
 131
 132 deListComp (BindStmt pat list locn : quals) cons nil
 133   = dsExpr list             `thenDs` \ list' ->
 134     let
 135         pat_ty      = outPatType pat
 136         nil_ty      = coreExprType nil
 137     in
 138     newSysLocalsDs [pat_ty, nil_ty]                     `thenDs` \ [x,ys] ->
 139
 140     dsListComp quals cons (Var ys)                      `thenDs` \ rest ->
 141     matchSimply (Var x) ListCompMatch pat
 142                 rest (Var ys)                           `thenDs` \ core_match ->
 143     bindNonRecDs (mkLams [x,ys] fn_body)                $ \ fn ->
 144     dsListExpr list (Var fn) nil
 145
 146
 147 data FExpr = FEOther CoreExpr                   -- Default case
 148            | FECons                             -- cons
 149            | FEConsComposedWith CoreExpr        -- (cons . e)
 150            | FENil                              -- nil
 151
 152 feComposeWith FECons g
 153   = returnDs (FEConsComposedWith g)
 154
 155 feComposeWith (FEOther f) g
 156   = composeWith f f     `thenDs` \ h ->
 157     returnDs (FEOther h)
 158
 159 feComposeWith (FEConsComposedWith f) g
 160   = composeWith f f     `thenDs` \ h ->
 161     returnDs (FEConsComposedWith h)
 162
 163
 164 composeWith f g
 165   = newSysLocalDs arg_ty        `thenDs` \ x ->
 166     returnDs (Lam x (App e (App f (Var x))))
 167   where
 168     arg_ty = case splitFunTy_maybe (coreExprType g) of
 169                 Just (arg_ty,_) -> arg_ty
 170                 other           -> panic "feComposeWith"
 171
 172 deListExpr :: TypecheckedHsExpr
 173            -> FExpr -> FExpr    -- Cons and nil expressions
 174            -> DsM CoreExpr
 175
 176 deListExpr cons nil (HsDoOut ListComp stmts _ _ _ result_ty src_loc)
 177   = deListComp stmts cons nil
 178
 179 deListExpr cons nil (HsVar map, _, [f,xs])
 180  | goodInst var mapIdKey = dsExpr f                     `thenDs` \ f' ->
 181                            feComposeWith cons f'        `thenDs` \ cons' ->
 182                            in
 183                            deListExpr xs cons' nil
 184
 185
 186 data HsExprForm = GoodForm What [Type] [TypecheckedHsExpr]
 187                 | BadForm
 188
 189 data What = HsMap | HsConcat | HsFilter |  HsZip | HsFoldr
 190
 191 analyseListProducer (HsVar v) ty_args val_args
 192   | good_inst mapIdKey    2 = GoodForm HsMap ty_args val_args
 193   | good_inst concatIdKey 1 = GoodForm HsConcat ty_args val_args
 194   | good_inst filterIdKey 2 = GoodForm HsFilter ty_args val_args
 195   | good_id   zipIdKey    2 = GoodForm HsZip    ty_args val_args
 196   | otherwise               =
 197   where
 198     good_inst key arity = isInstIdOf key v   && result_is_list && n_args == arity
 199     good_id   key arity = getUnique v == key && result_is_list && n_args == arity
 200
 201     n_args :: Int
 202     n_args = length val_args
 203
 204     result_is_list = resultTyIsList (idType v) ty_args val_args
 205
 206 resultTyIsList ty ty_args val_args
 207   = go ty ty_args
 208   where
 209     go1 ty (_:tys) = case splitForAllTy_maybe ty of
 210                         Just (_,ty) -> go1 ty tys
 211                         Nothing     -> False
 212     go1 ty [] = go2 ty val_args
 213
 214     go2 ty (_:args) = case splitFunTy_maybe of
 215                         Just (_,ty) -> go2 ty args
 216                         Nothing     -> False
 217
 218     go2 ty [] = case splitTyConApp_maybe of
 219                   Just (tycon, [_]) | tycon == listTyCon -> True
 220                   other                                  -> False
 221
 222
 223 \begin{code}
 224 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
 225
 226 deListComp [ReturnStmt expr] list               -- Figure 7.4, SLPJ, p 135, rule C above
 227   = dsExpr expr                 `thenDs` \ core_expr ->
 228     returnDs (mkConApp consDataCon [Type (coreExprType core_expr), core_expr, list])
 229
 230 deListComp (GuardStmt guard locn : quals) list  -- rule B above
 231   = dsExpr guard                `thenDs` \ core_guard ->
 232     deListComp quals list       `thenDs` \ core_rest ->
 233     returnDs (mkIfThenElse core_guard core_rest list)
 234
 235 -- [e | let B, qs] = let B in [e | qs]
 236 deListComp (LetStmt binds : quals) list
 237   = deListComp quals list       `thenDs` \ core_rest ->
 238     dsLet binds core_rest
 239
 240 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
 241   = dsExpr list1                    `thenDs` \ core_list1 ->
 242     let
 243         u3_ty@u1_ty = coreExprType core_list1   -- two names, same thing
 244
 245         -- u1_ty is a [alpha] type, and u2_ty = alpha
 246         u2_ty = outPatType pat
 247
 248         res_ty = coreExprType core_list2
 249         h_ty   = u1_ty `mkFunTy` res_ty
 250     in
 251     newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]  `thenDs` \ [h, u1, u2, u3] ->
 252
 253     -- the "fail" value ...
 254     let
 255         core_fail   = App (Var h) (Var u3)
 256         letrec_body = App (Var h) core_list1
 257     in
 258     deListComp quals core_fail                  `thenDs` \ rest_expr ->
 259     matchSimply (Var u2) ListCompMatch pat
 260                 rest_expr core_fail             `thenDs` \ core_match ->
 261     let
 262         rhs = Lam u1 $
 263               Case (Var u1) u1 [(DataCon nilDataCon,  [],       core_list2),
 264                                 (DataCon consDataCon, [u2, u3], core_match)]
 265     in
 266     returnDs (Let (Rec [(h, rhs)]) letrec_body)
 267 \end{code}
 268
 269 %************************************************************************
 270 %*                                                                      *
 271 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
 272 %*                                                                      *
 273 %************************************************************************
 274
 275 @dfListComp@ are the rules used with foldr/build turned on:
 276 \begin{verbatim}
 277 TE < [ e | ] >>          c n = c e n
 278 TE << [ e | b , q ] >>   c n = if b then TE << [ e | q ] >> c n else n
 279 TE << [ e | p <- l , q ] c n =  foldr
 280                         (\ TE << p >> b -> TE << [ e | q ] >> c b
 281                            _          b  -> b)  n l
 282 \end{verbatim}
 283 \begin{code}
 284 dfListComp :: Type -> Id                -- 'c'; its type and id
 285            -> Type -> Id                -- 'n'; its type and id
 286            -> [TypecheckedStmt]         -- the rest of the qual's
 287            -> DsM CoreExpr
 288
 289 dfListComp c_ty c_id n_ty n_id [ReturnStmt expr]
 290   = dsExpr expr                 `thenDs` \ core_expr ->
 291     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 292
 293 dfListComp c_ty c_id n_ty n_id (GuardStmt guard locn  : quals)
 294   = dsExpr guard                                `thenDs` \ core_guard ->
 295     dfListComp c_ty c_id n_ty n_id quals        `thenDs` \ core_rest ->
 296     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
 297
 298 dfListComp c_ty c_id n_ty n_id (LetStmt binds : quals)
 299   -- new in 1.3, local bindings
 300   = dfListComp c_ty c_id n_ty n_id quals        `thenDs` \ core_rest ->
 301     dsLet binds core_rest
 302
 303 dfListComp c_ty c_id n_ty n_id (BindStmt pat list1 locn : quals)
 304     -- evaluate the two lists
 305   = dsExpr list1                                `thenDs` \ core_list1 ->
 306
 307     -- find the required type
 308
 309     let p_ty   = outPatType pat
 310         b_ty   = n_ty           -- alias b_ty to n_ty
 311         fn_ty  = mkFunTys [p_ty, b_ty] b_ty
 312         lst_ty = coreExprType core_list1
 313     in
 314
 315     -- create some new local id's
 316
 317     newSysLocalsDs [b_ty,p_ty,fn_ty,lst_ty]             `thenDs` \ [b,p,fn,lst] ->
 318
 319     -- build rest of the comprehesion
 320
 321     dfListComp c_ty c_id b_ty b quals                   `thenDs` \ core_rest ->
 322     -- build the pattern match
 323
 324     matchSimply (Var p) ListCompMatch pat core_rest (Var b)     `thenDs` \ core_expr ->
 325
 326     -- now build the outermost foldr, and return
 327
 328     returnDs (
 329       mkLets
 330         [NonRec fn (mkLams [p, b] core_expr),
 331          NonRec lst core_list1]
 332         (mkFoldr p_ty n_ty fn n_id lst)
 333     )
 334 \end{code}
 335
 336
 337 @mkBuild@ is sugar for building a build!
 338
 339 @mkbuild ty tv c n e@ $Rightarrow$ @build ty (/\ tv -> \ c n -> e)@
 340 @ty@ is the type of the list.
 341 @tv@ is always a new type variable.
 342 @c,n@ are Id's for the abstract cons and nil, @g@ for let binding the argument argument.
 343         c :: a -> b -> b
 344         n :: b
 345         v :: (\/ b . (a -> b -> b) -> b -> b) -> [a]
 346 --  \/ a .  (\/ b . (a -> b -> b) -> b -> b) -> [a]
 347 @e@ is the object right inside the @build@
 348
 349 \begin{code}
 350 mkBuild :: Type
 351         -> TyVar
 352         -> Id
 353         -> Id
 354         -> Id
 355         -> CoreExpr -- template
 356         -> CoreExpr -- template
 357
 358 mkBuild ty tv c n g expr
 359   = Let (NonRec g (mkLams [tv, c,n] expr))
 360         (mkApps (Var buildId) [Type ty, Var g])
 361
 362 buildId = error "DsListComp: buildId"
 363
 364 mkFoldr a b f z xs
 365   = mkApps (mkTyApps (Var foldrId) [a,b]) [Var f, Var z, Var xs]
 366 \end{code}
 367