{-# LANGUAGE CPP, NamedFieldPuns #-}
{-# LANGUAGE TypeFamilies #-}
module DsListComp ( dsListComp, dsPArrComp, dsMonadComp ) where
#include "HsVersions.h"
import GhcPrelude
import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLExprNoLP, dsLocalBinds, dsSyntaxExpr )
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 TcType
import ListSetOps( getNth )
import Util
dsListComp :: [ExprLStmt GhcTc]
-> 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 GhcTc GhcTc) -> DsM (CoreExpr, Type)
dsInnerListComp (ParStmtBlock stmts bndrs _)
= do { let bndrs_tuple_type = mkBigCoreVarTupTy bndrs
list_ty = mkListTy bndrs_tuple_type
; expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTupId bndrs)]) list_ty
; return (expr, bndrs_tuple_type) }
dsTransStmt :: ExprStmt GhcTc -> DsM (CoreExpr, LPat GhcTc)
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
let 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' ]
dsNoLevPoly (tcFunResultTyN (length usingArgs') (exprType usingExpr'))
(text "In the result of a" <+> quotes (text "using") <+> text "function:" <+> ppr using)
let pat = mkBigLHsVarPatTupId to_bndrs
return (bound_unzipped_inner_list_expr', pat)
dsTransStmt _ = panic "dsTransStmt: Not given a TransStmt"
deListComp :: [ExprStmt GhcTc] -> 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 <- dsLExprNoLP 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 = mkBigLHsPatTupId pats
pats = map mkBigLHsVarPatTupId bndrs_s
deListComp (RecStmt {} : _) _ = panic "deListComp RecStmt"
deListComp (ApplicativeStmt {} : _) _ =
panic "deListComp ApplicativeStmt"
deBindComp :: OutPat GhcTc
-> CoreExpr
-> [ExprStmt GhcTc]
-> CoreExpr
-> DsM (Expr Id)
deBindComp pat core_list1 quals core_list2 = do
let u3_ty@u1_ty = exprType core_list1
let u2_ty = hsLPatType pat
let 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 GhcTc]
-> DsM CoreExpr
dfListComp _ _ [] = panic "dfListComp"
dfListComp c_id n_id (LastStmt body _ _ : quals)
= ASSERT( null quals )
do { core_body <- dsLExprNoLP 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"
dfListComp _ _ (ApplicativeStmt {} : _) =
panic "dfListComp ApplicativeStmt"
dfBindComp :: Id -> Id
-> (LPat GhcTc, CoreExpr)
-> [ExprStmt GhcTc]
-> DsM CoreExpr
dfBindComp c_id n_id (pat, core_list1) quals = do
let x_ty = hsLPatType pat
let b_ty = idType n_id
b <- newSysLocalDs b_ty
x <- newSysLocalDs 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 GhcTc]
-> DsM CoreExpr
dsPArrComp (ParStmt qss _ _ _ : quals) = dePArrParComp qss quals
dsPArrComp (BindStmt p e _ _ _ : qs) = do
filterP <- dsDPHBuiltin filterPVar
ce <- dsLExprNoLP 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 GhcTc]
-> LPat GhcTc
-> 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 lds@(L _ ds) : qs) pa cea = do
mapP <- dsDPHBuiltin mapPVar
let xs = collectLocalBinders ds
ty'cea = parrElemType cea
v <- newSysLocalDs ty'cea
clet <- dsLocalBinds lds (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 = text "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"
dePArrComp (ApplicativeStmt {} : _) _ _ =
panic "DsListComp.dePArrComp: ApplicativeStmt"
dePArrParComp :: [ParStmtBlock GhcTc GhcTc] -> [ExprStmt GhcTc] -> 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 GhcTc
-> LHsExpr GhcTc
-> DsM (CoreExpr, Type)
deLambda ty p e =
mkLambda ty p =<< dsLExpr e
mkLambda :: Type
-> LPat GhcTc
-> CoreExpr
-> DsM (CoreExpr, Type)
mkLambda ty p ce = do
v <- newSysLocalDs ty
let errMsg = text "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 GhcTc] -> DsM CoreExpr
dsMonadComp stmts = dsMcStmts stmts
dsMcStmts :: [ExprLStmt GhcTc] -> DsM CoreExpr
dsMcStmts [] = panic "dsMcStmts"
dsMcStmts (L loc stmt : lstmts) = putSrcSpanDs loc (dsMcStmt stmt lstmts)
dsMcStmt :: ExprStmt GhcTc -> [ExprLStmt GhcTc] -> DsM CoreExpr
dsMcStmt (LastStmt body _ ret_op) stmts
= ASSERT( null stmts )
do { body' <- dsLExpr body
; dsSyntaxExpr ret_op [body'] }
dsMcStmt (LetStmt binds) stmts
= do { rest <- dsMcStmts stmts
; dsLocalBinds binds rest }
dsMcStmt (BindStmt pat rhs bind_op fail_op bind_ty) stmts
= do { rhs' <- dsLExpr rhs
; dsMcBindStmt pat rhs' bind_op fail_op bind_ty stmts }
dsMcStmt (BodyStmt exp then_exp guard_exp _) stmts
= do { exp' <- dsLExpr exp
; rest <- dsMcStmts stmts
; guard_exp' <- dsSyntaxExpr guard_exp [exp']
; dsSyntaxExpr then_exp [guard_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_bind_arg_ty = n_tup_ty'
, trS_fmap = fmap_op, trS_form = form }) stmts_rest
= do { let (from_bndrs, to_bndrs) = unzip bndrs
; let 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'] }
; let tup_n_ty' = mkBigCoreVarTupTy to_bndrs
; body <- dsMcStmts stmts_rest
; n_tup_var' <- newSysLocalDsNoLP 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'
; dsSyntaxExpr bind_op [rhs', Lam n_tup_var' body'] }
dsMcStmt (ParStmt blocks mzip_op bind_op bind_ty) stmts_rest
= do { exps_w_tys <- mapM ds_inner blocks
; mzip_op' <- dsExpr mzip_op
; let
pats = [ mkBigLHsVarPatTupId 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 bind_ty 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 GhcTc
-> CoreExpr
-> SyntaxExpr GhcTc
-> SyntaxExpr GhcTc
-> Type
-> [ExprLStmt GhcTc]
-> DsM CoreExpr
dsMcBindStmt pat rhs' bind_op fail_op res1_ty stmts
= do { body <- dsMcStmts stmts
; var <- selectSimpleMatchVarL pat
; match <- matchSinglePat (Var var) (StmtCtxt DoExpr) pat
res1_ty (cantFailMatchResult body)
; match_code <- handle_failure pat match fail_op
; dsSyntaxExpr bind_op [rhs', Lam var match_code] }
where
handle_failure pat match fail_op
| matchCanFail match
= do { dflags <- getDynFlags
; fail_msg <- mkStringExpr (mk_fail_msg dflags pat)
; fail_expr <- dsSyntaxExpr fail_op [fail_msg]
; extractMatchResult match fail_expr }
| 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 GhcTc]
-> [Id]
-> SyntaxExpr GhcTc
-> DsM CoreExpr
dsInnerMonadComp stmts bndrs ret_op
= dsMcStmts (stmts ++ [noLoc (LastStmt (mkBigLHsVarTupId bndrs) False ret_op)])
mkMcUnzipM :: TransForm
-> HsExpr GhcTcId
-> 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])) }