DsListComp.lhs 16 KB
Newer Older
1
%
2
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3
%
chak's avatar
chak committed
4
\section[DsListComp]{Desugaring list comprehensions and array comprehensions}
5 6

\begin{code}
chak's avatar
chak committed
7
module DsListComp ( dsListComp, dsPArrComp ) where
8

9 10
#include "HsVersions.h"

11
import {-# SOURCE #-} DsExpr ( dsExpr, dsLet )
12

13
import BasicTypes	( Boxity(..) )
chak's avatar
chak committed
14
import TyCon		( tyConName )
15
import HsSyn		( Pat(..), HsExpr(..), Stmt(..),
16
			  HsMatchContext(..), HsStmtContext(..),
17
			  collectHsBinders )
chak's avatar
chak committed
18
import TcHsSyn		( TypecheckedStmt, TypecheckedPat, TypecheckedHsExpr,
19
			  hsPatType )
20
import CoreSyn
21

22
import DsMonad		-- the monadery used in the desugarer
23
import DsUtils
24

25
import CmdLineOpts	( DynFlag(..), dopt, opt_RulesOff )
26
import CoreUtils	( exprType, mkIfThenElse )
27
import Id		( idType )
28
import Var              ( Id )
chak's avatar
chak committed
29 30
import Type		( mkTyVarTy, mkFunTys, mkFunTy, Type,
			  splitTyConApp_maybe )
31
import TysPrim		( alphaTyVar )
32
import TysWiredIn	( nilDataCon, consDataCon, trueDataConId, falseDataConId, 
33
			  unitDataConId, unitTy, mkListTy )
34
import Match		( matchSimply )
35 36
import PrelNames	( foldrName, buildName, replicatePName, mapPName, 
			  filterPName, zipPName, crossPName, parrTyConName ) 
chak's avatar
chak committed
37
import PrelInfo		( pAT_ERROR_ID )
38
import SrcLoc		( noSrcLoc )
chak's avatar
chak committed
39
import Panic		( panic )
40 41 42 43 44 45 46 47 48
\end{code}

List comprehensions may be desugared in one of two ways: ``ordinary''
(as you would expect if you read SLPJ's book) and ``with foldr/build
turned on'' (if you read Gill {\em et al.}'s paper on the subject).

There will be at least one ``qualifier'' in the input.

\begin{code}
49 50 51 52 53
dsListComp :: [TypecheckedStmt] 
	   -> Type		-- Type of list elements
	   -> DsM CoreExpr

dsListComp quals elt_ty
54 55 56 57 58 59 60 61 62
  = getDOptsDs  `thenDs` \dflags ->
    if opt_RulesOff || dopt Opt_IgnoreInterfacePragmas dflags
	-- Either rules are switched off, or we are ignoring what there are;
	-- Either way foldr/build won't happen, so use the more efficient
	-- Wadler-style desugaring
  	|| isParallelComp quals
		-- Foldr-style desugaring can't handle
		-- parallel list comprehensions
  	then deListComp quals (mkNilExpr elt_ty)
63

64 65 66
   else		-- Foldr/build should be enabled, so desugar 
		-- into foldrs and builds
    newTyVarsDs [alphaTyVar]    `thenDs` \ [n_tyvar] ->
67 68 69
    let
	n_ty = mkTyVarTy n_tyvar
        c_ty = mkFunTys [elt_ty, n_ty] n_ty
70
    in
71 72
    newSysLocalsDs [c_ty,n_ty]		`thenDs` \ [c, n] ->
    dfListComp c n quals		`thenDs` \ result ->
73
    dsLookupGlobalId buildName	`thenDs` \ build_id ->
74 75
    returnDs (Var build_id `App` Type elt_ty 
			   `App` mkLams [n_tyvar, c, n] result)
76

77 78
  where isParallelComp (ParStmt bndrstmtss : _) = True
	isParallelComp _                        = False
79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
\end{code}

%************************************************************************
%*									*
\subsection[DsListComp-ordinary]{Ordinary desugaring of list comprehensions}
%*									*
%************************************************************************

Just as in Phil's chapter~7 in SLPJ, using the rules for
optimally-compiled list comprehensions.  This is what Kevin followed
as well, and I quite happily do the same.  The TQ translation scheme
transforms a list of qualifiers (either boolean expressions or
generators) into a single expression which implements the list
comprehension.  Because we are generating 2nd-order polymorphic
lambda-calculus, calls to NIL and CONS must be applied to a type
argument, as well as their usual value arguments.
\begin{verbatim}
TE << [ e | qs ] >>  =  TQ << [ e | qs ] ++ Nil (typeOf e) >>

(Rule C)
TQ << [ e | ] ++ L >> = Cons (typeOf e) TE <<e>> TE <<L>>

(Rule B)
TQ << [ e | b , qs ] ++ L >> =
    if TE << b >> then TQ << [ e | qs ] ++ L >> else TE << L >>

(Rule A')
TQ << [ e | p <- L1, qs ]  ++  L2 >> =
  letrec
    h = \ u1 ->
    	  case u1 of
	    []        ->  TE << L2 >>
	    (u2 : u3) ->
		  (( \ TE << p >> -> ( TQ << [e | qs]  ++  (h u3) >> )) u2)
		    [] (h u3)
  in
    h ( TE << L1 >> )

"h", "u1", "u2", and "u3" are new variables.
\end{verbatim}

@deListComp@ is the TQ translation scheme.  Roughly speaking, @dsExpr@
is the TE translation scheme.  Note that we carry around the @L@ list
already desugared.  @dsListComp@ does the top TE rule mentioned above.

124 125 126 127 128
To the above, we add an additional rule to deal with parallel list
comprehensions.  The translation goes roughly as follows:
     [ e | p1 <- e11, let v1 = e12, p2 <- e13
         | q1 <- e21, let v2 = e22, q2 <- e23]
     =>
129 130 131 132 133 134
     [ e | ((x1, .., xn), (y1, ..., ym)) <-
               zip [(x1,..,xn) | p1 <- e11, let v1 = e12, p2 <- e13]
                   [(y1,..,ym) | q1 <- e21, let v2 = e22, q2 <- e23]]
where (x1, .., xn) are the variables bound in p1, v1, p2
      (y1, .., ym) are the variables bound in q1, v2, q2

135
In the translation below, the ParStmt branch translates each parallel branch
136 137 138 139 140 141 142 143
into a sub-comprehension, and desugars each independently.  The resulting lists
are fed to a zip function, we create a binding for all the variables bound in all
the comprehensions, and then we hand things off the the desugarer for bindings.
The zip function is generated here a) because it's small, and b) because then we
don't have to deal with arbitrary limits on the number of zip functions in the
prelude, nor which library the zip function came from.
The introduced tuples are Boxed, but only because I couldn't get it to work
with the Unboxed variety.
144

145
\begin{code}
146

147
deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
148

149 150
deListComp (ParStmt stmtss_w_bndrs : quals) list
  = mapDs do_list_comp stmtss_w_bndrs	`thenDs` \ exps ->
151 152 153 154 155 156
    mkZipBind qual_tys			`thenDs` \ (zip_fn, zip_rhs) ->

	-- Deal with [e | pat <- zip l1 .. ln] in example above
    deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps)) 
		   quals list

157 158 159 160 161 162
  where 
	bndrs_s = map snd stmtss_w_bndrs

	-- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
	pat	 = TuplePat pats Boxed
	pats	 = map mk_hs_tuple_pat bndrs_s
163 164

	-- Types of (x1,..,xn), (y1,..,yn) etc
165
	qual_tys = map mk_bndrs_tys bndrs_s
166

167
	do_list_comp (stmts, bndrs)
168 169 170
	  = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
		       (mk_bndrs_tys bndrs)

171
	mk_bndrs_tys bndrs = mkCoreTupTy (map idType bndrs)
172

173
	-- Last: the one to return
174
deListComp [ResultStmt expr locn] list	-- Figure 7.4, SLPJ, p 135, rule C above
175
  = dsExpr expr			`thenDs` \ core_expr ->
176
    returnDs (mkConsExpr (exprType core_expr) core_expr list)
177

178
	-- Non-last: must be a guard
179
deListComp (ExprStmt guard ty locn : quals) list	-- rule B above
180 181
  = dsExpr guard       		`thenDs` \ core_guard ->
    deListComp quals list	`thenDs` \ core_rest ->
182
    returnDs (mkIfThenElse core_guard core_rest list)
183

184
-- [e | let B, qs] = let B in [e | qs]
185
deListComp (LetStmt binds : quals) list
186 187
  = deListComp quals list	`thenDs` \ core_rest ->
    dsLet binds core_rest
188

189
deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
190
  = dsExpr list1		    `thenDs` \ core_list1 ->
191
    deBindComp pat core_list1 quals core_list2
192 193
\end{code}

194

195
\begin{code}
196 197
deBindComp pat core_list1 quals core_list2
  = let
198
	u3_ty@u1_ty = exprType core_list1	-- two names, same thing
199 200

	-- u1_ty is a [alpha] type, and u2_ty = alpha
201
	u2_ty = hsPatType pat
202

203
	res_ty = exprType core_list2
204
	h_ty   = u1_ty `mkFunTy` res_ty
205
    in
206
    newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]	`thenDs` \ [h, u1, u2, u3] ->
207

208
    -- the "fail" value ...
209 210 211 212
    let
	core_fail   = App (Var h) (Var u3)
	letrec_body = App (Var h) core_list1
    in
213
    deListComp quals core_fail			`thenDs` \ rest_expr ->
214
    matchSimply (Var u2) (StmtCtxt ListComp) pat
215
		rest_expr core_fail		`thenDs` \ core_match ->
216 217
    let
	rhs = Lam u1 $
218 219
	      Case (Var u1) u1 [(DataAlt nilDataCon,  [], 	core_list2),
			        (DataAlt consDataCon, [u2, u3], core_match)]
220 221
    in
    returnDs (Let (Rec [(h, rhs)]) letrec_body)
222 223
\end{code}

224

225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
\begin{code}
mkZipBind :: [Type] -> DsM (Id, CoreExpr)
-- mkZipBind [t1, t2] 
-- = (zip, \as1:[t1] as2:[t2] 
--	   -> case as1 of 
--		[] -> []
--		(a1:as'1) -> case as2 of
--				[] -> []
--				(a2:as'2) -> (a2,a2) : zip as'1 as'2)]

mkZipBind elt_tys 
  = mapDs newSysLocalDs  list_tys	`thenDs` \ ass ->
    mapDs newSysLocalDs  elt_tys	`thenDs` \ as' ->
    mapDs newSysLocalDs  list_tys	`thenDs` \ as's ->
    newSysLocalDs zip_fn_ty		`thenDs` \ zip_fn ->
    let 
241 242 243
	inner_rhs = mkConsExpr ret_elt_ty 
			(mkCoreTup (map Var as'))
			(mkVarApps (Var zip_fn) as's)
244 245 246 247 248
	zip_body  = foldr mk_case inner_rhs (zip3 ass as' as's)
    in
    returnDs (zip_fn, mkLams ass zip_body)
  where
    list_tys   = map mkListTy elt_tys
249
    ret_elt_ty = mkCoreTupTy elt_tys
250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267
    zip_fn_ty  = mkFunTys list_tys (mkListTy ret_elt_ty)

    mk_case (as, a', as') rest
	  = Case (Var as) as [(DataAlt nilDataCon,  [],        mkNilExpr ret_elt_ty),
			      (DataAlt consDataCon, [a', as'], rest)]

-- Helper functions that makes an HsTuple only for non-1-sized tuples
mk_hs_tuple_expr :: [Id] -> TypecheckedHsExpr
mk_hs_tuple_expr []   = HsVar unitDataConId
mk_hs_tuple_expr [id] = HsVar id
mk_hs_tuple_expr ids  = ExplicitTuple [ HsVar i | i <- ids ] Boxed

mk_hs_tuple_pat :: [Id] -> TypecheckedPat
mk_hs_tuple_pat [b] = VarPat b
mk_hs_tuple_pat bs  = TuplePat (map VarPat bs) Boxed
\end{code}


268 269 270 271 272 273 274
%************************************************************************
%*									*
\subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
%*									*
%************************************************************************

@dfListComp@ are the rules used with foldr/build turned on:
275

276
\begin{verbatim}
277 278 279 280 281 282 283 284
TE[ e | ]            c n = c e n
TE[ e | b , q ]      c n = if b then TE[ e | q ] c n else n
TE[ e | p <- l , q ] c n = let 
				f = \ x b -> case x of
						  p -> TE[ e | q ] c b
						  _ -> b
			   in
			   foldr f n l
285
\end{verbatim}
286

287
\begin{code}
288
dfListComp :: Id -> Id			-- 'c' and 'n'
289
	   -> [TypecheckedStmt] 	-- the rest of the qual's
290
	   -> DsM CoreExpr
291

292
	-- Last: the one to return
293
dfListComp c_id n_id [ResultStmt expr locn]
294
  = dsExpr expr			`thenDs` \ core_expr ->
295
    returnDs (mkApps (Var c_id) [core_expr, Var n_id])
296

297
	-- Non-last: must be a guard
298
dfListComp c_id n_id (ExprStmt guard ty locn  : quals)
299
  = dsExpr guard               			`thenDs` \ core_guard ->
300
    dfListComp c_id n_id quals	`thenDs` \ core_rest ->
301
    returnDs (mkIfThenElse core_guard core_rest (Var n_id))
302

303
dfListComp c_id n_id (LetStmt binds : quals)
304
  -- new in 1.3, local bindings
305
  = dfListComp c_id n_id quals	`thenDs` \ core_rest ->
306
    dsLet binds core_rest
307

308
dfListComp c_id n_id (BindStmt pat list1 locn : quals)
309 310 311 312
    -- evaluate the two lists
  = dsExpr list1				`thenDs` \ core_list1 ->

    -- find the required type
313
    let x_ty   = hsPatType pat
314
	b_ty   = idType n_id
315 316 317
    in

    -- create some new local id's
318
    newSysLocalsDs [b_ty,x_ty]			`thenDs` \ [b,x] ->
319 320

    -- build rest of the comprehesion
321
    dfListComp c_id b quals			`thenDs` \ core_rest ->
322 323

    -- build the pattern match
324
    matchSimply (Var x) (StmtCtxt ListComp) 
325
		pat core_rest (Var b)		`thenDs` \ core_expr ->
326 327

    -- now build the outermost foldr, and return
328
    dsLookupGlobalId foldrName		`thenDs` \ foldr_id ->
329
    returnDs (
330 331 332 333 334
      Var foldr_id `App` Type x_ty 
		   `App` Type b_ty
		   `App` mkLams [x, b] core_expr
		   `App` Var n_id
		   `App` core_list1
335
    )
336 337
\end{code}

chak's avatar
chak committed
338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353
%************************************************************************
%*									*
\subsection[DsPArrComp]{Desugaring of array comprehensions}
%*									*
%************************************************************************

\begin{code}

-- entry point for desugaring a parallel array comprehension
--
--   [:e | qss:] = <<[:e | qss:]>> () [:():]
--
dsPArrComp      :: [TypecheckedStmt] 
	        -> Type		    -- Don't use; called with `undefined' below
	        -> DsM CoreExpr
dsPArrComp qs _  =
354
  dsLookupGlobalId replicatePName			  `thenDs` \repP ->
chak's avatar
chak committed
355
  let unitArray = mkApps (Var repP) [Type unitTy, 
356
				     mkIntExpr 1, 
357
				     mkCoreTup []]
chak's avatar
chak committed
358 359
  in
  dePArrComp qs (TuplePat [] Boxed) unitArray
360

chak's avatar
chak committed
361 362 363 364 365 366 367 368 369 370
-- the work horse
--
dePArrComp :: [TypecheckedStmt] 
	   -> TypecheckedPat		-- the current generator pattern
	   -> CoreExpr			-- the current generator expression
	   -> DsM CoreExpr
--
--  <<[:e' | :]>> pa ea = mapP (\pa -> e') ea
--
dePArrComp [ResultStmt e' _] pa cea =
371
  dsLookupGlobalId mapPName				  `thenDs` \mapP    ->
chak's avatar
chak committed
372 373 374 375 376 377 378 379 380
  let ty = parrElemType cea
  in
  deLambda ty pa e'					  `thenDs` \(clam, 
								     ty'e') ->
  returnDs $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea]
--
--  <<[:e' | b, qs:]>> pa ea = <<[:e' | qs:]>> pa (filterP (\pa -> b) ea)
--
dePArrComp (ExprStmt b _ _ : qs) pa cea =
381
  dsLookupGlobalId filterPName			  `thenDs` \filterP  ->
chak's avatar
chak committed
382 383 384 385 386 387 388 389 390 391 392
  let ty = parrElemType cea
  in
  deLambda ty pa b					  `thenDs` \(clam,_) ->
  dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
--
--  <<[:e' | p <- e, qs:]>> pa ea = 
--    let ef = filterP (\x -> case x of {p -> True; _ -> False}) e
--    in
--    <<[:e' | qs:]>> (pa, p) (crossP ea ef)
--
dePArrComp (BindStmt p e _ : qs) pa cea =
393 394
  dsLookupGlobalId filterPName			  `thenDs` \filterP ->
  dsLookupGlobalId crossPName			  `thenDs` \crossP  ->
395
  dsExpr e					  `thenDs` \ce      ->
chak's avatar
chak committed
396 397
  let ty'cea = parrElemType cea
      ty'ce  = parrElemType ce
398 399
      false  = Var falseDataConId
      true   = Var trueDataConId
chak's avatar
chak committed
400 401
  in
  newSysLocalDs ty'ce					  `thenDs` \v       ->
402
  matchSimply (Var v) (StmtCtxt PArrComp) p true false      `thenDs` \pred    ->
chak's avatar
chak committed
403 404 405 406 407 408 409 410 411 412 413 414 415
  let cef    = mkApps (Var filterP) [Type ty'ce, mkLams [v] pred, ce]
      ty'cef = ty'ce				-- filterP preserves the type
      pa'    = TuplePat [pa, p] Boxed
  in
  dePArrComp qs pa' (mkApps (Var crossP) [Type ty'cea, Type ty'cef, cea, cef])
--
--  <<[:e' | let ds, qs:]>> pa ea = 
--    <<[:e' | qs:]>> (pa, (x_1, ..., x_n)) 
--		      (mapP (\v@pa -> (v, let ds in (x_1, ..., x_n))) ea)
--  where
--    {x_1, ..., x_n} = DV (ds)		-- Defined Variables
--
dePArrComp (LetStmt ds : qs) pa cea =
416 417
  dsLookupGlobalId mapPName				  `thenDs` \mapP    ->
  let xs     = collectHsBinders ds
chak's avatar
chak committed
418 419 420
      ty'cea = parrElemType cea
  in
  newSysLocalDs ty'cea					  `thenDs` \v       ->
421
  dsLet ds (mkCoreTup (map Var xs))			  `thenDs` \clet    ->
chak's avatar
chak committed
422
  newSysLocalDs (exprType clet)				  `thenDs` \let'v   ->
423 424
  let projBody = mkDsLet (NonRec let'v clet) $ 
		 mkCoreTup [Var v, Var let'v]
chak's avatar
chak committed
425 426 427 428
      errTy    = exprType projBody
      errMsg   = "DsListComp.dePArrComp: internal error!"
  in
  mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
429
  matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr  `thenDs` \ccase   ->
chak's avatar
chak committed
430 431 432 433 434 435 436 437 438 439 440
  let pa'    = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
      proj   = mkLams [v] ccase
  in
  dePArrComp qs pa' (mkApps (Var mapP) [Type ty'cea, proj, cea])
--
--  <<[:e' | qs | qss:]>> pa ea = 
--    <<[:e' | qss:]>> (pa, (x_1, ..., x_n)) 
--		       (zipP ea <<[:(x_1, ..., x_n) | qs:]>>)
--    where
--      {x_1, ..., x_n} = DV (qs)
--
441 442
dePArrComp (ParStmt []             : qss2) pa cea = dePArrComp qss2 pa cea
dePArrComp (ParStmt ((qs, xs):qss) : qss2) pa cea =
443
  dsLookupGlobalId zipPName				  `thenDs` \zipP    ->
chak's avatar
chak committed
444 445 446 447 448 449 450 451
  let pa'     = TuplePat [pa, TuplePat (map VarPat xs) Boxed] Boxed
      ty'cea  = parrElemType cea
      resStmt = ResultStmt (ExplicitTuple (map HsVar xs) Boxed) noSrcLoc
  in
  dsPArrComp (qs ++ [resStmt]) undefined		  `thenDs` \cqs     ->
  let ty'cqs = parrElemType cqs
      cea'   = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
  in
452
  dePArrComp (ParStmt qss : qss2) pa' cea'
chak's avatar
chak committed
453 454 455 456 457 458 459 460 461 462 463 464 465 466

-- generate Core corresponding to `\p -> e'
--
deLambda        :: Type			-- type of the argument
		-> TypecheckedPat	-- argument pattern
		-> TypecheckedHsExpr	-- body
		-> DsM (CoreExpr, Type)
deLambda ty p e  =
  newSysLocalDs ty					  `thenDs` \v       ->
  dsExpr e						  `thenDs` \ce      ->
  let errTy    = exprType ce
      errMsg   = "DsListComp.deLambda: internal error!"
  in
  mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
467
  matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr	  `thenDs` \res	    ->
chak's avatar
chak committed
468 469 470 471 472 473 474 475 476 477 478 479
  returnDs (mkLams [v] res, errTy)

-- obtain the element type of the parallel array produced by the given Core
-- expression
--
parrElemType   :: CoreExpr -> Type
parrElemType e  = 
  case splitTyConApp_maybe (exprType e) of
    Just (tycon, [ty]) | tyConName tycon == parrTyConName -> ty
    _							  -> panic
      "DsListComp.parrElemType: not a parallel array type"
\end{code}