Add kind equalities to GHC.
[ghc.git] / compiler / vectorise / Vectorise / Utils / Closure.hs
1 -- |Utils concerning closure construction and application.
2
3 module Vectorise.Utils.Closure
4 ( mkClosure
5 , mkClosureApp
6 , buildClosures
7 )
8 where
9
10 import Vectorise.Builtins
11 import Vectorise.Vect
12 import Vectorise.Monad
13 import Vectorise.Utils.Base
14 import Vectorise.Utils.PADict
15 import Vectorise.Utils.Hoisting
16
17 import CoreSyn
18 import Type
19 import MkCore
20 import CoreUtils
21 import TyCon
22 import DataCon
23 import MkId
24 import TysWiredIn
25 import BasicTypes( Boxity(..) )
26 import FastString
27
28
29 -- |Make a closure.
30 --
31 mkClosure :: Type -- ^ Type of the argument.
32 -> Type -- ^ Type of the result.
33 -> Type -- ^ Type of the environment.
34 -> VExpr -- ^ The function to apply.
35 -> VExpr -- ^ The environment to use.
36 -> VM VExpr
37 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
38 = do dict <- paDictOfType env_ty
39 mkv <- builtin closureVar
40 mkl <- builtin liftedClosureVar
41 return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
42 Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
43
44 -- |Make a closure application.
45 --
46 mkClosureApp :: Type -- ^ Type of the argument.
47 -> Type -- ^ Type of the result.
48 -> VExpr -- ^ Closure to apply.
49 -> VExpr -- ^ Argument to use.
50 -> VM VExpr
51 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
52 = do vapply <- builtin applyVar
53 lapply <- builtin liftedApplyVar
54 lc <- builtin liftingContext
55 return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
56 Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
57
58 -- |Build a set of 'n' closures corresponding to an 'n'-ary vectorised function. The length of
59 -- the list of types of arguments determines the arity.
60 --
61 -- In addition to a set of type variables, a set of value variables is passed during closure
62 -- /construction/. In contrast, the closure environment and the arguments are passed during closure
63 -- application.
64 --
65 buildClosures :: [TyVar] -- ^ Type variables passed during closure construction.
66 -> [Var] -- ^ Variables passed during closure construction.
67 -> [VVar] -- ^ Variables in the environment.
68 -> [Type] -- ^ Type of the arguments.
69 -> Type -- ^ Type of result.
70 -> VM VExpr
71 -> VM VExpr
72 buildClosures _tvs _vars _env [] _res_ty mk_body
73 = mk_body
74 buildClosures tvs vars env [arg_ty] res_ty mk_body
75 = buildClosure tvs vars env arg_ty res_ty mk_body
76 buildClosures tvs vars env (arg_ty : arg_tys) res_ty mk_body
77 = do { res_ty' <- mkClosureTypes arg_tys res_ty
78 ; arg <- newLocalVVar (fsLit "x") arg_ty
79 ; buildClosure tvs vars env arg_ty res_ty'
80 . hoistPolyVExpr tvs vars (Inline (length env + 1))
81 $ do { lc <- builtin liftingContext
82 ; clo <- buildClosures tvs vars (env ++ [arg]) arg_tys res_ty mk_body
83 ; return $ vLams lc (env ++ [arg]) clo
84 }
85 }
86
87 -- Build a closure taking one extra argument during closure application.
88 --
89 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
90 -- where
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
93 --
94 -- In addition to a set of type variables, a set of value variables is passed during closure
95 -- /construction/. In contrast, the closure environment and the closure argument are passed during
96 -- closure application.
97 --
98 buildClosure :: [TyVar] -- ^Type variables passed during closure construction.
99 -> [Var] -- ^Variables passed during closure construction.
100 -> [VVar] -- ^Variables in the environment.
101 -> Type -- ^Type of the closure argument.
102 -> Type -- ^Type of the result.
103 -> VM VExpr
104 -> VM VExpr
105 buildClosure tvs vars vvars arg_ty res_ty mk_body
106 = do { (env_ty, env, bind) <- buildEnv vvars
107 ; env_bndr <- newLocalVVar (fsLit "env") env_ty
108 ; arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
109
110 -- generate the closure function as a hoisted binding
111 ; fn <- hoistPolyVExpr tvs vars (Inline 2) $
112 do { lc <- builtin liftingContext
113 ; body <- mk_body
114 ; return . vLams lc [env_bndr, arg_bndr]
115 $ bind (vVar env_bndr)
116 (vVarApps lc body (vvars ++ [arg_bndr]))
117 }
118
119 ; mkClosure arg_ty res_ty env_ty fn env
120 }
121
122 -- Build the environment for a single closure.
123 --
124 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
125 buildEnv []
126 = do
127 ty <- voidType
128 void <- builtin voidVar
129 pvoid <- builtin pvoidVar
130 return (ty, vVar (void, pvoid), \_ body -> body)
131 buildEnv [v]
132 = return (vVarType v, vVar v,
133 \env body -> vLet (vNonRec v env) body)
134 buildEnv vs
135 = do (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
136
137 let venv_con = tupleDataCon Boxed (length vs)
138 [lenv_con] = tyConDataCons lenv_tc
139
140 venv = mkCoreTup (map Var vvs)
141 lenv = Var (dataConWrapId lenv_con)
142 `mkTyApps` lenv_tyargs
143 `mkApps` map Var lvs
144
145 vbind env body = mkWildCase env ty (exprType body)
146 [(DataAlt venv_con, vvs, body)]
147
148 lbind env body =
149 let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
150 in
151 mkWildCase scrut (exprType scrut) (exprType body)
152 [(DataAlt lenv_con, lvs, body)]
153
154 bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
155 lbind lenv lbody)
156
157 return (ty, (venv, lenv), bind)
158 where
159 (vvs, lvs) = unzip vs
160 tys = map vVarType vs
161 ty = mkBoxedTupleTy tys