[project @ 2003-06-02 13:28:08 by simonpj]
[ghc-hetmet.git] / ghc / compiler / deSugar / DsListComp.lhs
index 431fb93..7af59eb 100644 (file)
@@ -1,19 +1,22 @@
 %
 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
 %
-\section[DsListComp]{Desugaring list comprehensions}
+\section[DsListComp]{Desugaring list comprehensions and array comprehensions}
 
 \begin{code}
-module DsListComp ( dsListComp ) where
+module DsListComp ( dsListComp, dsPArrComp ) where
 
 #include "HsVersions.h"
 
 import {-# SOURCE #-} DsExpr ( dsExpr, dsLet )
 
 import BasicTypes      ( Boxity(..) )
-import HsSyn           ( OutPat(..), HsExpr(..), Stmt(..), HsMatchContext(..) )
-import TcHsSyn         ( TypecheckedStmt )
-import DsHsSyn         ( outPatType )
+import TyCon           ( tyConName )
+import HsSyn           ( Pat(..), HsExpr(..), Stmt(..),
+                         HsMatchContext(..), HsStmtContext(..),
+                         collectHsBinders )
+import TcHsSyn         ( TypecheckedStmt, TypecheckedPat, TypecheckedHsExpr,
+                         hsPatType )
 import CoreSyn
 
 import DsMonad         -- the monadery used in the desugarer
@@ -23,13 +26,18 @@ import CmdLineOpts  ( opt_FoldrBuildOn )
 import CoreUtils       ( exprType, mkIfThenElse )
 import Id              ( idType )
 import Var              ( Id )
-import Type            ( mkTyVarTy, mkFunTys, mkFunTy, Type )
+import Type            ( mkTyVarTy, mkFunTys, mkFunTy, Type,
+                         splitTyConApp_maybe )
 import TysPrim         ( alphaTyVar )
-import TysWiredIn      ( nilDataCon, consDataCon, unitDataConId, tupleCon, mkListTy, mkTupleTy )
+import TysWiredIn      ( nilDataCon, consDataCon, trueDataConId, falseDataConId, 
+                         unitDataConId, unitTy,
+                         mkListTy, mkTupleTy )
 import Match           ( matchSimply )
-import PrelNames       ( foldrName, buildName )
+import PrelNames       ( foldrName, buildName, replicatePName, mapPName, 
+                         filterPName, zipPName, crossPName, parrTyConName ) 
+import PrelInfo                ( pAT_ERROR_ID )
 import SrcLoc          ( noSrcLoc )
-import List            ( zip4 )
+import Panic           ( panic )
 \end{code}
 
 List comprehensions may be desugared in one of two ways: ``ordinary''
@@ -44,7 +52,8 @@ dsListComp :: [TypecheckedStmt]
           -> DsM CoreExpr
 
 dsListComp quals elt_ty
-  | not opt_FoldrBuildOn                -- Be boring
+  |  not opt_FoldrBuildOn               -- Be boring
+  || isParallelComp quals
   = deListComp quals (mkNilExpr elt_ty)
 
   | otherwise                           -- foldr/build lives!
@@ -55,9 +64,12 @@ dsListComp quals elt_ty
     in
     newSysLocalsDs [c_ty,n_ty]         `thenDs` \ [c, n] ->
     dfListComp c n quals               `thenDs` \ result ->
-    dsLookupGlobalValue buildName      `thenDs` \ build_id ->
+    dsLookupGlobalId buildName `thenDs` \ build_id ->
     returnDs (Var build_id `App` Type elt_ty 
                           `App` mkLams [n_tyvar, c, n] result)
+
+  where isParallelComp (ParStmtOut bndrstmtss : _) = True
+       isParallelComp _                           = False
 \end{code}
 
 %************************************************************************
@@ -108,9 +120,12 @@ comprehensions.  The translation goes roughly as follows:
      [ e | p1 <- e11, let v1 = e12, p2 <- e13
          | q1 <- e21, let v2 = e22, q2 <- e23]
      =>
-     [ e | ((p1,v1,p2), (q1,v2,q2)) <-
-               zip [(p1,v1,p2) | p1 <- e11, let v1 = e12, p2 <- e13]
-                   [(q1,v2,q2) | q1 <- e21, let v2 = e22, q2 <- e23]]
+     [ 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
+
 In the translation below, the ParStmtOut branch translates each parallel branch
 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
@@ -126,64 +141,33 @@ with the Unboxed variety.
 deListComp :: [TypecheckedStmt] -> CoreExpr -> DsM CoreExpr
 
 deListComp (ParStmtOut bndrstmtss : quals) list
-  = mapDs doListComp qualss    `thenDs` \ exps ->
-    mapDs genAS  bndrss                `thenDs` \ ass ->
-    mapDs genA   bndrss                `thenDs` \ as ->
-    mapDs genAS' bndrss                `thenDs` \ as's ->
-    let retTy = myTupleTy Boxed (length bndrss) qualTys
-       zipTy = foldr mkFunTy (mkListTy retTy) (map mkListTy qualTys)
-    in
-    newSysLocalDs zipTy                `thenDs` \ zipFn ->
-    let target = mkConsExpr retTy (mkTupleExpr as) (foldl App (Var zipFn) (map Var as's))
-       zipExp = mkLet zipFn (zip4 bndrss ass as as's) exps target
-    in
-    deBindComp pat zipExp quals list
-  where (bndrss, stmtss) = unzip bndrstmtss
-       pats = map (\ps -> mkTuplePat (map VarPat ps)) bndrss
-       mkTuplePat [p] = p
-       mkTuplePat ps  = TuplePat ps Boxed
-       pat  = TuplePat pats Boxed
-
-       qualss = map mkQuals bndrstmtss
-       mkQuals (bndrs, stmts) = (bndrs, stmts ++ [ExprStmt (myTupleExpr bndrs) noSrcLoc])
-
-       qualTys = map mkBndrsTy bndrss
-       mkBndrsTy bndrs = myTupleTy Boxed (length bndrs) (map idType bndrs)
-
-       doListComp (bndrs, stmts)
-         = dsListComp stmts (mkBndrsTy bndrs)
-       genA   bndrs = newSysLocalDs (mkBndrsTy bndrs)
-       genAS  bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
-       genAS' bndrs = newSysLocalDs (mkListTy (mkBndrsTy bndrs))
-
-       mkLet zipFn vars exps target
-         = Let (Rec [(zipFn,
-                      foldr Lam (mkBody target vars) (map getAs vars))])
-               (foldl App (Var zipFn) exps)
-       getAs (_, as, _, _) = as
-       mkBody target vars
-         = foldr mkCase (foldr mkTuplCase target vars) vars
-       mkCase (ps, as, a, as') rest
-         = Case (Var as) as [(DataAlt nilDataCon, [], mkConApp nilDataCon []),
-                             (DataAlt consDataCon, [a, as'], rest)]
-       mkTuplCase ([p], as, a, as') rest
-         = App (Lam p rest) (Var a)
-       mkTuplCase (ps, as, a, as') rest
-         = Case (Var a) a [(DataAlt (tupleCon Boxed (length ps)), ps, rest)]
-
-       myTupleTy boxity arity [ty] = ty
-       myTupleTy boxity arity tys  = mkTupleTy boxity arity tys
-       myTupleExpr []   = HsVar unitDataConId
-       myTupleExpr [id] = HsVar id
-       myTupleExpr ids  = ExplicitTuple [ HsVar i | i <- ids ] Boxed
+  = mapDs do_list_comp bndrstmtss      `thenDs` \ exps ->
+    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
+
+  where -- pat is the pattern ((x1,..,xn), (y1,..,ym)) in the example above
+       pat            = TuplePat pats Boxed
+       pats           = map (\(bs,_) -> mk_hs_tuple_pat bs) bndrstmtss
+
+       -- Types of (x1,..,xn), (y1,..,yn) etc
+       qual_tys = [ mk_bndrs_tys bndrs | (bndrs,_) <- bndrstmtss ]
+
+       do_list_comp (bndrs, stmts)
+         = dsListComp (stmts ++ [ResultStmt (mk_hs_tuple_expr bndrs) noSrcLoc])
+                      (mk_bndrs_tys bndrs)
+
+       mk_bndrs_tys bndrs = mk_tuple_ty (map idType bndrs)
 
        -- Last: the one to return
-deListComp [ExprStmt expr locn] list   -- Figure 7.4, SLPJ, p 135, rule C above
+deListComp [ResultStmt expr locn] list -- Figure 7.4, SLPJ, p 135, rule C above
   = dsExpr expr                        `thenDs` \ core_expr ->
     returnDs (mkConsExpr (exprType core_expr) core_expr list)
 
        -- Non-last: must be a guard
-deListComp (ExprStmt guard locn : quals) list  -- rule B above
+deListComp (ExprStmt guard ty locn : quals) list       -- rule B above
   = dsExpr guard                       `thenDs` \ core_guard ->
     deListComp quals list      `thenDs` \ core_rest ->
     returnDs (mkIfThenElse core_guard core_rest list)
@@ -196,13 +180,16 @@ deListComp (LetStmt binds : quals) list
 deListComp (BindStmt pat list1 locn : quals) core_list2 -- rule A' above
   = dsExpr list1                   `thenDs` \ core_list1 ->
     deBindComp pat core_list1 quals core_list2
+\end{code}
 
+
+\begin{code}
 deBindComp pat core_list1 quals core_list2
   = let
        u3_ty@u1_ty = exprType core_list1       -- two names, same thing
 
        -- u1_ty is a [alpha] type, and u2_ty = alpha
-       u2_ty = outPatType pat
+       u2_ty = hsPatType pat
 
        res_ty = exprType core_list2
        h_ty   = u1_ty `mkFunTy` res_ty
@@ -215,7 +202,7 @@ deBindComp pat core_list1 quals core_list2
        letrec_body = App (Var h) core_list1
     in
     deListComp quals core_fail                 `thenDs` \ rest_expr ->
-    matchSimply (Var u2) ListComp pat
+    matchSimply (Var u2) (StmtCtxt ListComp) pat
                rest_expr core_fail             `thenDs` \ core_match ->
     let
        rhs = Lam u1 $
@@ -226,6 +213,54 @@ deBindComp pat core_list1 quals core_list2
 \end{code}
 
 
+\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 
+       inner_rhs = mkConsExpr ret_elt_ty 
+                       (mkCoreTup (map Var as'))
+                       (mkVarApps (Var zip_fn) as's)
+       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
+    ret_elt_ty = mk_tuple_ty elt_tys
+    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 function 
+mk_tuple_ty :: [Type] -> Type
+mk_tuple_ty [ty] = ty
+mk_tuple_ty tys  = mkTupleTy Boxed (length tys) tys
+
+-- 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}
+
+
 %************************************************************************
 %*                                                                     *
 \subsection[DsListComp-foldr-build]{Foldr/Build desugaring of list comprehensions}
@@ -251,12 +286,12 @@ dfListComp :: Id -> Id                    -- 'c' and 'n'
           -> DsM CoreExpr
 
        -- Last: the one to return
-dfListComp c_id n_id [ExprStmt expr locn]
+dfListComp c_id n_id [ResultStmt expr locn]
   = dsExpr expr                        `thenDs` \ core_expr ->
     returnDs (mkApps (Var c_id) [core_expr, Var n_id])
 
        -- Non-last: must be a guard
-dfListComp c_id n_id (ExprStmt guard locn  : quals)
+dfListComp c_id n_id (ExprStmt guard ty locn  : quals)
   = dsExpr guard                                       `thenDs` \ core_guard ->
     dfListComp c_id n_id quals `thenDs` \ core_rest ->
     returnDs (mkIfThenElse core_guard core_rest (Var n_id))
@@ -271,7 +306,7 @@ dfListComp c_id n_id (BindStmt pat list1 locn : quals)
   = dsExpr list1                               `thenDs` \ core_list1 ->
 
     -- find the required type
-    let x_ty   = outPatType pat
+    let x_ty   = hsPatType pat
        b_ty   = idType n_id
     in
 
@@ -282,10 +317,11 @@ dfListComp c_id n_id (BindStmt pat list1 locn : quals)
     dfListComp c_id b quals                    `thenDs` \ core_rest ->
 
     -- build the pattern match
-    matchSimply (Var x) ListComp pat core_rest (Var b) `thenDs` \ core_expr ->
+    matchSimply (Var x) (StmtCtxt ListComp) 
+               pat core_rest (Var b)           `thenDs` \ core_expr ->
 
     -- now build the outermost foldr, and return
-    dsLookupGlobalValue foldrName              `thenDs` \ foldr_id ->
+    dsLookupGlobalId foldrName         `thenDs` \ foldr_id ->
     returnDs (
       Var foldr_id `App` Type x_ty 
                   `App` Type b_ty
@@ -295,4 +331,145 @@ dfListComp c_id n_id (BindStmt pat list1 locn : quals)
     )
 \end{code}
 
+%************************************************************************
+%*                                                                     *
+\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 _  =
+  dsLookupGlobalId replicatePName                        `thenDs` \repP ->
+  let unitArray = mkApps (Var repP) [Type unitTy, 
+                                    mkIntExpr 1, 
+                                    mkCoreTup []]
+  in
+  dePArrComp qs (TuplePat [] Boxed) unitArray
+
+-- 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 =
+  dsLookupGlobalId mapPName                              `thenDs` \mapP    ->
+  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 =
+  dsLookupGlobalId filterPName                   `thenDs` \filterP  ->
+  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 =
+  dsLookupGlobalId filterPName                   `thenDs` \filterP ->
+  dsLookupGlobalId crossPName                    `thenDs` \crossP  ->
+  dsExpr e                                       `thenDs` \ce      ->
+  let ty'cea = parrElemType cea
+      ty'ce  = parrElemType ce
+      false  = Var falseDataConId
+      true   = Var trueDataConId
+  in
+  newSysLocalDs ty'ce                                    `thenDs` \v       ->
+  matchSimply (Var v) (StmtCtxt PArrComp) p true false      `thenDs` \pred    ->
+  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 =
+  dsLookupGlobalId mapPName                              `thenDs` \mapP    ->
+  let xs     = collectHsBinders ds
+      ty'cea = parrElemType cea
+  in
+  newSysLocalDs ty'cea                                   `thenDs` \v       ->
+  dsLet ds (mkCoreTup (map Var xs))                      `thenDs` \clet    ->
+  newSysLocalDs (exprType clet)                                  `thenDs` \let'v   ->
+  let projBody = mkDsLet (NonRec let'v clet) $ 
+                mkCoreTup [Var v, Var let'v]
+      errTy    = exprType projBody
+      errMsg   = "DsListComp.dePArrComp: internal error!"
+  in
+  mkErrorAppDs pAT_ERROR_ID errTy errMsg                  `thenDs` \cerr    ->
+  matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr  `thenDs` \ccase   ->
+  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)
+--
+dePArrComp (ParStmtOut []             : qss2) pa cea = dePArrComp qss2 pa cea
+dePArrComp (ParStmtOut ((xs, qs):qss) : qss2) pa cea =
+  dsLookupGlobalId zipPName                              `thenDs` \zipP    ->
+  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
+  dePArrComp (ParStmtOut qss : qss2) pa' cea'
+
+-- 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    ->
+  matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr      `thenDs` \res     ->
+  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}