module DsListComp ( dsListComp, dsMonadComp ) where
#include "HsVersions.h"
import GhcPrelude
import 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) }
dsInnerListComp (XParStmtBlock{}) = panic "dsInnerListComp"
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 noExt 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"
deListComp (XStmtLR {} : _) _ =
panic "deListComp XStmtLR"
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"
dfListComp _ _ (XStmtLR {} : _) =
panic "dfListComp XStmtLR"
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
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 bind_ty pat rhs bind_op fail_op) 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_ext = 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 bind_ty blocks mzip_op bind_op) 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) }
ds_inner (XParStmtBlock{}) = panic "dsMcStmt"
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 <- matchSinglePatVar 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 noExt (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])) }