2 -- | Utils concerning closure construction and application.
3 module Vectorise.Utils.Closure (
11 import Vectorise.Builtins
13 import Vectorise.Monad
14 import Vectorise.Utils.Base
15 import Vectorise.Utils.PADict
16 import Vectorise.Utils.Hoisting
27 import BasicTypes( Boxity(..) )
33 :: Type -- ^ Type of the argument.
34 -> Type -- ^ Type of the result.
35 -> Type -- ^ Type of the environment.
36 -> VExpr -- ^ The function to apply.
37 -> VExpr -- ^ The environment to use.
40 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
41 = do Just dict <- paDictOfType env_ty
42 mkv <- builtin closureVar
43 mkl <- builtin liftedClosureVar
44 return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
45 Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
48 -- | Make a closure application.
50 :: Type -- ^ Type of the argument.
51 -> Type -- ^ Type of the result.
52 -> VExpr -- ^ Closure to apply.
53 -> VExpr -- ^ Argument to use.
56 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
57 = do vapply <- builtin applyVar
58 lapply <- builtin liftedApplyVar
59 lc <- builtin liftingContext
60 return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
61 Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
68 -> [Type] -- ^ Type of the arguments.
69 -> Type -- ^ Type of result.
73 buildClosures _ _ [] _ mk_body
76 buildClosures tvs vars [arg_ty] res_ty mk_body
77 = buildClosure tvs vars arg_ty res_ty mk_body
79 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
80 = do res_ty' <- mkClosureTypes arg_tys res_ty
81 arg <- newLocalVVar (fsLit "x") arg_ty
82 buildClosure tvs vars arg_ty res_ty'
83 . hoistPolyVExpr tvs (Inline (length vars + 1))
85 lc <- builtin liftingContext
86 clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
87 return $ vLams lc (vars ++ [arg]) clo
90 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
92 -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
93 -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
95 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
96 buildClosure tvs vars arg_ty res_ty mk_body
98 (env_ty, env, bind) <- buildEnv vars
99 env_bndr <- newLocalVVar (fsLit "env") env_ty
100 arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
102 fn <- hoistPolyVExpr tvs (Inline 2)
104 lc <- builtin liftingContext
106 return . vLams lc [env_bndr, arg_bndr]
107 $ bind (vVar env_bndr)
108 (vVarApps lc body (vars ++ [arg_bndr]))
110 mkClosure arg_ty res_ty env_ty fn env
113 -- Environments ---------------------------------------------------------------
114 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
118 void <- builtin voidVar
119 pvoid <- builtin pvoidVar
120 return (ty, vVar (void, pvoid), \_ body -> body)
122 buildEnv [v] = return (vVarType v, vVar v,
123 \env body -> vLet (vNonRec v env) body)
126 = do (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
128 let venv_con = tupleCon Boxed (length vs)
129 [lenv_con] = tyConDataCons lenv_tc
131 venv = mkCoreTup (map Var vvs)
132 lenv = Var (dataConWrapId lenv_con)
133 `mkTyApps` lenv_tyargs
136 vbind env body = mkWildCase env ty (exprType body)
137 [(DataAlt venv_con, vvs, body)]
140 let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
142 mkWildCase scrut (exprType scrut) (exprType body)
143 [(DataAlt lenv_con, lvs, body)]
145 bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
148 return (ty, (venv, lenv), bind)
150 (vvs, lvs) = unzip vs
151 tys = map vVarType vs
152 ty = mkBoxedTupleTy tys