Commit 6326f92d authored by rl@cse.unsw.edu.au's avatar rl@cse.unsw.edu.au
Browse files

Improve closure generation for functions with multiple parameters

parent 76cec9c6
module VectUtils (
collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
collectAnnValBinders,
splitClosureTy,
mkPADictType, mkPArrayType,
paDictArgType, paDictOfType,
......@@ -7,7 +8,7 @@ module VectUtils (
polyAbstract, polyApply, polyVApply,
lookupPArrayFamInst,
hoistExpr, hoistPolyVExpr, takeHoisted,
buildClosure
buildClosure, buildClosures
) where
#include "HsVersions.h"
......@@ -46,6 +47,12 @@ collectAnnTypeBinders expr = go [] expr
go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
go bs e = (reverse bs, e)
collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
collectAnnValBinders expr = go [] expr
where
go bs (_, AnnLam b e) | isId b = go (b:bs) e
go bs e = (reverse bs, e)
isAnnTypeArg :: AnnExpr b ann -> Bool
isAnnTypeArg (_, AnnType t) = True
isAnnTypeArg _ = False
......@@ -72,6 +79,20 @@ splitPArrayTy ty
| otherwise = pprPanic "splitPArrayTy" (ppr ty)
mkClosureType :: Type -> Type -> VM Type
mkClosureType arg_ty res_ty
= do
tc <- builtin closureTyCon
return $ mkTyConApp tc [arg_ty, res_ty]
mkClosureTypes :: [Type] -> Type -> VM Type
mkClosureTypes arg_tys res_ty
= do
tc <- builtin closureTyCon
return $ foldr (mk tc) res_ty arg_tys
where
mk tc arg_ty res_ty = mkTyConApp tc [arg_ty, res_ty]
mkPADictType :: Type -> VM Type
mkPADictType ty
= do
......@@ -227,11 +248,24 @@ mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
buildClosures :: [TyVar] -> Var -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
buildClosures tvs lc vars [arg_ty] res_ty mk_body
= buildClosure tvs lc vars arg_ty res_ty mk_body
buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
= do
res_ty' <- mkClosureTypes arg_tys res_ty
arg <- newLocalVVar FSLIT("x") arg_ty
buildClosure tvs lc vars arg_ty res_ty'
. hoistPolyVExpr FSLIT("fn") tvs
$ do
clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body
return $ vLams lc (vars ++ [arg]) clo
-- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
-- where
-- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
-- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
--
buildClosure :: [TyVar] -> Var -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
buildClosure tvs lv vars arg_ty res_ty mk_body
= do
......
......@@ -242,26 +242,31 @@ vectExpr lc (_, AnnLet (AnnRec bs) body)
where
(bndrs, rhss) = unzip bs
vectExpr lc e@(_, AnnLam bndr body)
| isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
vectExpr lc e@(fvs, AnnLam bndr _)
| not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
| otherwise = vectLam lc fvs bs body
where
(bs,body) = collectAnnValBinders e
vectExpr lc (fvs, AnnLam bndr body)
vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam lc fvs bs body
= do
tyvars <- localTyVars
(vs, vvs) <- readLEnv $ \env ->
unzip [(var, vv) | var <- varSetElems fvs
, Just vv <- [lookupVarEnv (local_vars env) var]]
arg_ty <- vectType (idType bndr)
res_ty <- vectType (exprType $ deAnnotate body)
buildClosure tyvars lc vvs arg_ty res_ty
arg_tys <- mapM (vectType . idType) bs
res_ty <- vectType (exprType $ deAnnotate body)
buildClosures tyvars lc vvs arg_tys res_ty
. hoistPolyVExpr FSLIT("fn") tyvars
$ do
new_lc <- newLocalVar FSLIT("lc") intPrimTy
(vbndrs, vbody) <- vectBndrsIn (vs ++ [bndr])
(vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
(vectExpr new_lc body)
return $ vLams new_lc vbndrs vbody
vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment