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
26 import BasicTypes( Boxity(..) )
32 :: Type -- ^ Type of the argument.
33 -> Type -- ^ Type of the result.
34 -> Type -- ^ Type of the environment.
35 -> VExpr -- ^ The function to apply.
36 -> VExpr -- ^ The environment to use.
39 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
40 = do dict <- paDictOfType env_ty
41 mkv <- builtin closureVar
42 mkl <- builtin liftedClosureVar
43 return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
44 Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
47 -- | Make a closure application.
49 :: Type -- ^ Type of the argument.
50 -> Type -- ^ Type of the result.
51 -> VExpr -- ^ Closure to apply.
52 -> VExpr -- ^ Argument to use.
55 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
56 = do vapply <- builtin applyVar
57 lapply <- builtin liftedApplyVar
58 lc <- builtin liftingContext
59 return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
60 Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
67 -> [Type] -- ^ Type of the arguments.
68 -> Type -- ^ Type of result.
72 buildClosures _ _ [] _ mk_body
75 buildClosures tvs vars [arg_ty] res_ty mk_body
76 = buildClosure tvs vars arg_ty res_ty mk_body
78 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
79 = do res_ty' <- mkClosureTypes arg_tys res_ty
80 arg <- newLocalVVar (fsLit "x") arg_ty
81 buildClosure tvs vars arg_ty res_ty'
82 . hoistPolyVExpr tvs (Inline (length vars + 1))
84 lc <- builtin liftingContext
85 clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
86 return $ vLams lc (vars ++ [arg]) clo
89 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
91 -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
92 -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
94 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
95 buildClosure tvs vars arg_ty res_ty mk_body
97 (env_ty, env, bind) <- buildEnv vars
98 env_bndr <- newLocalVVar (fsLit "env") env_ty
99 arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
101 fn <- hoistPolyVExpr tvs (Inline 2)
103 lc <- builtin liftingContext
105 return . vLams lc [env_bndr, arg_bndr]
106 $ bind (vVar env_bndr)
107 (vVarApps lc body (vars ++ [arg_bndr]))
109 mkClosure arg_ty res_ty env_ty fn env
112 -- Environments ---------------------------------------------------------------
113 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
117 void <- builtin voidVar
118 pvoid <- builtin pvoidVar
119 return (ty, vVar (void, pvoid), \_ body -> body)
121 buildEnv [v] = return (vVarType v, vVar v,
122 \env body -> vLet (vNonRec v env) body)
125 = do (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
127 let venv_con = tupleCon Boxed (length vs)
128 [lenv_con] = tyConDataCons lenv_tc
130 venv = mkCoreTup (map Var vvs)
131 lenv = Var (dataConWrapId lenv_con)
132 `mkTyApps` lenv_tyargs
135 vbind env body = mkWildCase env ty (exprType body)
136 [(DataAlt venv_con, vvs, body)]
139 let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
141 mkWildCase scrut (exprType scrut) (exprType body)
142 [(DataAlt lenv_con, lvs, body)]
144 bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
147 return (ty, (venv, lenv), bind)
149 (vvs, lvs) = unzip vs
150 tys = map vVarType vs
151 ty = mkBoxedTupleTy tys