module DsListComp ( dsListComp, dsPArrComp, dsMonadComp ) where
#include "HsVersions.h"
import DsExpr ( dsExpr, dsLExpr, dsLocalBinds )
import HsSyn
import TcHsSyn
import CoreSyn
import MkCore
import DsMonad
import DsUtils
import DynFlags
import CoreUtils
import Id
import Type
import TysWiredIn
import Match
import PrelNames
import SrcLoc
import Outputable
import FastString
import TcType
import ListSetOps( getNth )
import Util
dsListComp :: [ExprLStmt Id]
-> Type
-> DsM CoreExpr
dsListComp lquals res_ty = do
dflags <- getDynFlags
let quals = map unLoc lquals
elt_ty = case tcTyConAppArgs res_ty of
[elt_ty] -> elt_ty
_ -> pprPanic "dsListComp" (ppr res_ty $$ ppr lquals)
if not (gopt Opt_EnableRewriteRules dflags) || gopt Opt_IgnoreInterfacePragmas dflags
|| isParallelComp quals
then deListComp quals (mkNilExpr elt_ty)
else mkBuildExpr elt_ty (\(c, _) (n, _) -> dfListComp c n quals)
where
isParallelComp = any isParallelStmt
isParallelStmt (ParStmt {}) = True
isParallelStmt _ = False
dsInnerListComp :: (ParStmtBlock Id Id) -> DsM (CoreExpr, Type)
dsInnerListComp (ParStmtBlock stmts bndrs _)
= do { expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTup bndrs)])
(mkListTy bndrs_tuple_type)
; return (expr, bndrs_tuple_type) }
where
bndrs_tuple_type = mkBigCoreVarTupTy bndrs
dsTransStmt :: ExprStmt Id -> DsM (CoreExpr, LPat Id)
dsTransStmt (TransStmt { trS_form = form, trS_stmts = stmts, trS_bndrs = binderMap
, trS_by = by, trS_using = using }) = do
let (from_bndrs, to_bndrs) = unzip binderMap
from_bndrs_tys = map idType from_bndrs
to_bndrs_tys = map idType to_bndrs
to_bndrs_tup_ty = mkBigCoreTupTy to_bndrs_tys
(expr, from_tup_ty) <- dsInnerListComp (ParStmtBlock stmts from_bndrs noSyntaxExpr)
usingExpr' <- dsLExpr using
usingArgs <- case by of
Nothing -> return [expr]
Just by_e -> do { by_e' <- dsLExpr by_e
; lam <- matchTuple from_bndrs by_e'
; return [lam, expr] }
unzip_stuff <- mkUnzipBind form from_bndrs_tys
map_id <- dsLookupGlobalId mapName
let
inner_list_expr = mkApps usingExpr' usingArgs
bound_unzipped_inner_list_expr
= case unzip_stuff of
Nothing -> inner_list_expr
Just (unzip_fn, unzip_rhs) -> Let (Rec [(unzip_fn, unzip_rhs)]) $
mkApps (Var map_id) $
[ Type (mkListTy from_tup_ty)
, Type to_bndrs_tup_ty
, Var unzip_fn
, inner_list_expr]
let pat = mkBigLHsVarPatTup to_bndrs
return (bound_unzipped_inner_list_expr, pat)
dsTransStmt _ = panic "dsTransStmt: Not given a TransStmt"
deListComp :: [ExprStmt Id] -> CoreExpr -> DsM CoreExpr
deListComp [] _ = panic "deListComp"
deListComp (LastStmt body _ : quals) list
=
ASSERT( null quals )
do { core_body <- dsLExpr body
; return (mkConsExpr (exprType core_body) core_body list) }
deListComp (BodyStmt guard _ _ _ : quals) list = do
core_guard <- dsLExpr guard
core_rest <- deListComp quals list
return (mkIfThenElse core_guard core_rest list)
deListComp (LetStmt binds : quals) list = do
core_rest <- deListComp quals list
dsLocalBinds binds core_rest
deListComp (stmt@(TransStmt {}) : quals) list = do
(inner_list_expr, pat) <- dsTransStmt stmt
deBindComp pat inner_list_expr quals list
deListComp (BindStmt pat list1 _ _ : quals) core_list2 = do
core_list1 <- dsLExpr list1
deBindComp pat core_list1 quals core_list2
deListComp (ParStmt stmtss_w_bndrs _ _ : quals) list
= do { exps_and_qual_tys <- mapM dsInnerListComp stmtss_w_bndrs
; let (exps, qual_tys) = unzip exps_and_qual_tys
; (zip_fn, zip_rhs) <- mkZipBind qual_tys
; deBindComp pat (Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps))
quals list }
where
bndrs_s = [bs | ParStmtBlock _ bs _ <- stmtss_w_bndrs]
pat = mkBigLHsPatTup pats
pats = map mkBigLHsVarPatTup bndrs_s
deListComp (RecStmt {} : _) _ = panic "deListComp RecStmt"
deBindComp :: OutPat Id
-> CoreExpr
-> [ExprStmt Id]
-> CoreExpr
-> DsM (Expr Id)
deBindComp pat core_list1 quals core_list2 = do
let
u3_ty@u1_ty = exprType core_list1
u2_ty = hsLPatType pat
res_ty = exprType core_list2
h_ty = u1_ty `mkFunTy` res_ty
[h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty]
let
core_fail = App (Var h) (Var u3)
letrec_body = App (Var h) core_list1
rest_expr <- deListComp quals core_fail
core_match <- matchSimply (Var u2) (StmtCtxt ListComp) pat rest_expr core_fail
let
rhs = Lam u1 $
Case (Var u1) u1 res_ty
[(DataAlt nilDataCon, [], core_list2),
(DataAlt consDataCon, [u2, u3], core_match)]
return (Let (Rec [(h, rhs)]) letrec_body)
dfListComp :: Id -> Id
-> [ExprStmt Id]
-> DsM CoreExpr
dfListComp _ _ [] = panic "dfListComp"
dfListComp c_id n_id (LastStmt body _ : quals)
= ASSERT( null quals )
do { core_body <- dsLExpr body
; return (mkApps (Var c_id) [core_body, Var n_id]) }
dfListComp c_id n_id (BodyStmt guard _ _ _ : quals) = do
core_guard <- dsLExpr guard
core_rest <- dfListComp c_id n_id quals
return (mkIfThenElse core_guard core_rest (Var n_id))
dfListComp c_id n_id (LetStmt binds : quals) = do
core_rest <- dfListComp c_id n_id quals
dsLocalBinds binds core_rest
dfListComp c_id n_id (stmt@(TransStmt {}) : quals) = do
(inner_list_expr, pat) <- dsTransStmt stmt
dfBindComp c_id n_id (pat, inner_list_expr) quals
dfListComp c_id n_id (BindStmt pat list1 _ _ : quals) = do
core_list1 <- dsLExpr list1
dfBindComp c_id n_id (pat, core_list1) quals
dfListComp _ _ (ParStmt {} : _) = panic "dfListComp ParStmt"
dfListComp _ _ (RecStmt {} : _) = panic "dfListComp RecStmt"
dfBindComp :: Id -> Id
-> (LPat Id, CoreExpr)
-> [ExprStmt Id]
-> DsM CoreExpr
dfBindComp c_id n_id (pat, core_list1) quals = do
let x_ty = hsLPatType pat
b_ty = idType n_id
[b, x] <- newSysLocalsDs [b_ty, x_ty]
core_rest <- dfListComp c_id b quals
core_expr <- matchSimply (Var x) (StmtCtxt ListComp)
pat core_rest (Var b)
mkFoldrExpr x_ty b_ty (mkLams [x, b] core_expr) (Var n_id) core_list1
mkZipBind :: [Type] -> DsM (Id, CoreExpr)
mkZipBind elt_tys = do
ass <- mapM newSysLocalDs elt_list_tys
as' <- mapM newSysLocalDs elt_tys
as's <- mapM newSysLocalDs elt_list_tys
zip_fn <- newSysLocalDs zip_fn_ty
let inner_rhs = mkConsExpr elt_tuple_ty
(mkBigCoreVarTup as')
(mkVarApps (Var zip_fn) as's)
zip_body = foldr mk_case inner_rhs (zip3 ass as' as's)
return (zip_fn, mkLams ass zip_body)
where
elt_list_tys = map mkListTy elt_tys
elt_tuple_ty = mkBigCoreTupTy elt_tys
elt_tuple_list_ty = mkListTy elt_tuple_ty
zip_fn_ty = mkFunTys elt_list_tys elt_tuple_list_ty
mk_case (as, a', as') rest
= Case (Var as) as elt_tuple_list_ty
[(DataAlt nilDataCon, [], mkNilExpr elt_tuple_ty),
(DataAlt consDataCon, [a', as'], rest)]
mkUnzipBind :: TransForm -> [Type] -> DsM (Maybe (Id, CoreExpr))
mkUnzipBind ThenForm _
= return Nothing
mkUnzipBind _ elt_tys
= do { ax <- newSysLocalDs elt_tuple_ty
; axs <- newSysLocalDs elt_list_tuple_ty
; ys <- newSysLocalDs elt_tuple_list_ty
; xs <- mapM newSysLocalDs elt_tys
; xss <- mapM newSysLocalDs elt_list_tys
; unzip_fn <- newSysLocalDs unzip_fn_ty
; [us1, us2] <- sequence [newUniqueSupply, newUniqueSupply]
; let nil_tuple = mkBigCoreTup (map mkNilExpr elt_tys)
concat_expressions = map mkConcatExpression (zip3 elt_tys (map Var xs) (map Var xss))
tupled_concat_expression = mkBigCoreTup concat_expressions
folder_body_inner_case = mkTupleCase us1 xss tupled_concat_expression axs (Var axs)
folder_body_outer_case = mkTupleCase us2 xs folder_body_inner_case ax (Var ax)
folder_body = mkLams [ax, axs] folder_body_outer_case
; unzip_body <- mkFoldrExpr elt_tuple_ty elt_list_tuple_ty folder_body nil_tuple (Var ys)
; return (Just (unzip_fn, mkLams [ys] unzip_body)) }
where
elt_tuple_ty = mkBigCoreTupTy elt_tys
elt_tuple_list_ty = mkListTy elt_tuple_ty
elt_list_tys = map mkListTy elt_tys
elt_list_tuple_ty = mkBigCoreTupTy elt_list_tys
unzip_fn_ty = elt_tuple_list_ty `mkFunTy` elt_list_tuple_ty
mkConcatExpression (list_element_ty, head, tail) = mkConsExpr list_element_ty head tail
dsPArrComp :: [ExprStmt Id]
-> DsM CoreExpr
dsPArrComp (ParStmt qss _ _ : quals) = dePArrParComp qss quals
dsPArrComp (BindStmt p e _ _ : qs) = do
filterP <- dsDPHBuiltin filterPVar
ce <- dsLExpr e
let ety'ce = parrElemType ce
false = Var falseDataConId
true = Var trueDataConId
v <- newSysLocalDs ety'ce
pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
let gen | isIrrefutableHsPat p = ce
| otherwise = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
dePArrComp qs p gen
dsPArrComp qs = do
sglP <- dsDPHBuiltin singletonPVar
let unitArray = mkApps (Var sglP) [Type unitTy, mkCoreTup []]
dePArrComp qs (noLoc $ WildPat unitTy) unitArray
dePArrComp :: [ExprStmt Id]
-> LPat Id
-> CoreExpr
-> DsM CoreExpr
dePArrComp [] _ _ = panic "dePArrComp"
dePArrComp (LastStmt e' _ : quals) pa cea
= ASSERT( null quals )
do { mapP <- dsDPHBuiltin mapPVar
; let ty = parrElemType cea
; (clam, ty'e') <- deLambda ty pa e'
; return $ mkApps (Var mapP) [Type ty, Type ty'e', clam, cea] }
dePArrComp (BodyStmt b _ _ _ : qs) pa cea = do
filterP <- dsDPHBuiltin filterPVar
let ty = parrElemType cea
(clam,_) <- deLambda ty pa b
dePArrComp qs pa (mkApps (Var filterP) [Type ty, clam, cea])
dePArrComp (BindStmt p e _ _ : qs) pa cea = do
filterP <- dsDPHBuiltin filterPVar
crossMapP <- dsDPHBuiltin crossMapPVar
ce <- dsLExpr e
let ety'cea = parrElemType cea
ety'ce = parrElemType ce
false = Var falseDataConId
true = Var trueDataConId
v <- newSysLocalDs ety'ce
pred <- matchSimply (Var v) (StmtCtxt PArrComp) p true false
let cef | isIrrefutableHsPat p = ce
| otherwise = mkApps (Var filterP) [Type ety'ce, mkLams [v] pred, ce]
(clam, _) <- mkLambda ety'cea pa cef
let ety'cef = ety'ce
pa' = mkLHsPatTup [pa, p]
dePArrComp qs pa' (mkApps (Var crossMapP)
[Type ety'cea, Type ety'cef, cea, clam])
dePArrComp (LetStmt ds : qs) pa cea = do
mapP <- dsDPHBuiltin mapPVar
let xs = collectLocalBinders ds
ty'cea = parrElemType cea
v <- newSysLocalDs ty'cea
clet <- dsLocalBinds ds (mkCoreTup (map Var xs))
let'v <- newSysLocalDs (exprType clet)
let projBody = mkCoreLet (NonRec let'v clet) $
mkCoreTup [Var v, Var let'v]
errTy = exprType projBody
errMsg = ptext (sLit "DsListComp.dePArrComp: internal error!")
cerr <- mkErrorAppDs pAT_ERROR_ID errTy errMsg
ccase <- matchSimply (Var v) (StmtCtxt PArrComp) pa projBody cerr
let pa' = mkLHsPatTup [pa, mkLHsPatTup (map nlVarPat xs)]
proj = mkLams [v] ccase
dePArrComp qs pa' (mkApps (Var mapP)
[Type ty'cea, Type errTy, proj, cea])
dePArrComp (ParStmt {} : _) _ _ =
panic "DsListComp.dePArrComp: malformed comprehension AST: ParStmt"
dePArrComp (TransStmt {} : _) _ _ = panic "DsListComp.dePArrComp: TransStmt"
dePArrComp (RecStmt {} : _) _ _ = panic "DsListComp.dePArrComp: RecStmt"
dePArrParComp :: [ParStmtBlock Id Id] -> [ExprStmt Id] -> DsM CoreExpr
dePArrParComp qss quals = do
(pQss, ceQss) <- deParStmt qss
dePArrComp quals pQss ceQss
where
deParStmt [] =
panic "DsListComp.dePArrComp: Empty parallel list comprehension"
deParStmt (ParStmtBlock qs xs _:qss) = do
let res_expr = mkLHsVarTuple xs
cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
parStmts qss (mkLHsVarPatTup xs) cqs
parStmts [] pa cea = return (pa, cea)
parStmts (ParStmtBlock qs xs _:qss) pa cea = do
zipP <- dsDPHBuiltin zipPVar
let pa' = mkLHsPatTup [pa, mkLHsVarPatTup xs]
ty'cea = parrElemType cea
res_expr = mkLHsVarTuple xs
cqs <- dsPArrComp (map unLoc qs ++ [mkLastStmt res_expr])
let ty'cqs = parrElemType cqs
cea' = mkApps (Var zipP) [Type ty'cea, Type ty'cqs, cea, cqs]
parStmts qss pa' cea'
deLambda :: Type
-> LPat Id
-> LHsExpr Id
-> DsM (CoreExpr, Type)
deLambda ty p e =
mkLambda ty p =<< dsLExpr e
mkLambda :: Type
-> LPat Id
-> CoreExpr
-> DsM (CoreExpr, Type)
mkLambda ty p ce = do
v <- newSysLocalDs ty
let errMsg = ptext (sLit "DsListComp.deLambda: internal error!")
ce'ty = exprType ce
cerr <- mkErrorAppDs pAT_ERROR_ID ce'ty errMsg
res <- matchSimply (Var v) (StmtCtxt PArrComp) p ce cerr
return (mkLams [v] res, ce'ty)
parrElemType :: CoreExpr -> Type
parrElemType e =
case splitTyConApp_maybe (exprType e) of
Just (tycon, [ty]) | tycon == parrTyCon -> ty
_ -> panic
"DsListComp.parrElemType: not a parallel array type"
dsMonadComp :: [ExprLStmt Id] -> DsM CoreExpr
dsMonadComp stmts = dsMcStmts stmts
dsMcStmts :: [ExprLStmt Id] -> DsM CoreExpr
dsMcStmts [] = panic "dsMcStmts"
dsMcStmts (L loc stmt : lstmts) = putSrcSpanDs loc (dsMcStmt stmt lstmts)
dsMcStmt :: ExprStmt Id -> [ExprLStmt Id] -> DsM CoreExpr
dsMcStmt (LastStmt body ret_op) stmts
= ASSERT( null stmts )
do { body' <- dsLExpr body
; ret_op' <- dsExpr ret_op
; return (App ret_op' body') }
dsMcStmt (LetStmt binds) stmts
= do { rest <- dsMcStmts stmts
; dsLocalBinds binds rest }
dsMcStmt (BindStmt pat rhs bind_op fail_op) stmts
= do { rhs' <- dsLExpr rhs
; dsMcBindStmt pat rhs' bind_op fail_op stmts }
dsMcStmt (BodyStmt exp then_exp guard_exp _) stmts
= do { exp' <- dsLExpr exp
; guard_exp' <- dsExpr guard_exp
; then_exp' <- dsExpr then_exp
; rest <- dsMcStmts stmts
; return $ mkApps then_exp' [ mkApps guard_exp' [exp']
, rest ] }
dsMcStmt (TransStmt { trS_stmts = stmts, trS_bndrs = bndrs
, trS_by = by, trS_using = using
, trS_ret = return_op, trS_bind = bind_op
, trS_fmap = fmap_op, trS_form = form }) stmts_rest
= do { let (from_bndrs, to_bndrs) = unzip bndrs
from_bndr_tys = map idType from_bndrs
; expr <- dsInnerMonadComp stmts from_bndrs return_op
; usingExpr' <- dsLExpr using
; usingArgs <- case by of
Nothing -> return [expr]
Just by_e -> do { by_e' <- dsLExpr by_e
; lam <- matchTuple from_bndrs by_e'
; return [lam, expr] }
; bind_op' <- dsExpr bind_op
; let bind_ty = exprType bind_op'
n_tup_ty = funArgTy $ funArgTy $ funResultTy bind_ty
tup_n_ty = mkBigCoreVarTupTy to_bndrs
; body <- dsMcStmts stmts_rest
; n_tup_var <- newSysLocalDs n_tup_ty
; tup_n_var <- newSysLocalDs tup_n_ty
; tup_n_expr <- mkMcUnzipM form fmap_op n_tup_var from_bndr_tys
; us <- newUniqueSupply
; let rhs' = mkApps usingExpr' usingArgs
body' = mkTupleCase us to_bndrs body tup_n_var tup_n_expr
; return (mkApps bind_op' [rhs', Lam n_tup_var body']) }
dsMcStmt (ParStmt blocks mzip_op bind_op) stmts_rest
= do { exps_w_tys <- mapM ds_inner blocks
; mzip_op' <- dsExpr mzip_op
; let
pats = [ mkBigLHsVarPatTup bs | ParStmtBlock _ bs _ <- blocks]
pat = foldr1 (\p1 p2 -> mkLHsPatTup [p1, p2]) pats
(rhs, _) = foldr1 (\(e1,t1) (e2,t2) ->
(mkApps mzip_op' [Type t1, Type t2, e1, e2],
mkBoxedTupleTy [t1,t2]))
exps_w_tys
; dsMcBindStmt pat rhs bind_op noSyntaxExpr stmts_rest }
where
ds_inner (ParStmtBlock stmts bndrs return_op)
= do { exp <- dsInnerMonadComp stmts bndrs return_op
; return (exp, mkBigCoreVarTupTy bndrs) }
dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
matchTuple ids body
= do { us <- newUniqueSupply
; tup_id <- newSysLocalDs (mkBigCoreVarTupTy ids)
; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
dsMcBindStmt :: LPat Id
-> CoreExpr
-> SyntaxExpr Id
-> SyntaxExpr Id
-> [ExprLStmt Id]
-> DsM CoreExpr
dsMcBindStmt pat rhs' bind_op fail_op stmts
= do { body <- dsMcStmts stmts
; bind_op' <- dsExpr bind_op
; var <- selectSimpleMatchVarL pat
; let bind_ty = exprType bind_op'
res1_ty = funResultTy (funArgTy (funResultTy bind_ty))
; match <- matchSinglePat (Var var) (StmtCtxt DoExpr) pat
res1_ty (cantFailMatchResult body)
; match_code <- handle_failure pat match fail_op
; return (mkApps bind_op' [rhs', Lam var match_code]) }
where
handle_failure pat match fail_op
| matchCanFail match
= do { fail_op' <- dsExpr fail_op
; dflags <- getDynFlags
; fail_msg <- mkStringExpr (mk_fail_msg dflags pat)
; extractMatchResult match (App fail_op' fail_msg) }
| otherwise
= extractMatchResult match (error "It can't fail")
mk_fail_msg :: DynFlags -> Located e -> String
mk_fail_msg dflags pat
= "Pattern match failure in monad comprehension at " ++
showPpr dflags (getLoc pat)
dsInnerMonadComp :: [ExprLStmt Id]
-> [Id]
-> HsExpr Id
-> DsM CoreExpr
dsInnerMonadComp stmts bndrs ret_op
= dsMcStmts (stmts ++ [noLoc (LastStmt (mkBigLHsVarTup bndrs) ret_op)])
mkMcUnzipM :: TransForm
-> SyntaxExpr TcId
-> Id
-> [Type]
-> DsM CoreExpr
mkMcUnzipM ThenForm _ ys _
= return (Var ys)
mkMcUnzipM _ fmap_op ys elt_tys
= do { fmap_op' <- dsExpr fmap_op
; xs <- mapM newSysLocalDs elt_tys
; let tup_ty = mkBigCoreTupTy elt_tys
; tup_xs <- newSysLocalDs tup_ty
; let mk_elt i = mkApps fmap_op'
[ Type tup_ty, Type (getNth elt_tys i)
, mk_sel i, Var ys]
mk_sel n = Lam tup_xs $
mkTupleSelector xs (getNth xs n) tup_xs (Var tup_xs)
; return (mkBigCoreTup (map mk_elt [0..length elt_tys 1])) }