a658253df778545bc3777bb2768103ec4ed3a90a
[ghc.git] / compiler / vectorise / VectMonad.hs
1 module VectMonad (
2 VM,
3
4 noV, tryV, maybeV, orElseV, localV, closedV, initV,
5 newLocalVar, newTyVar,
6
7 Builtins(..), paDictTyCon,
8 builtin,
9
10 GlobalEnv(..),
11 readGEnv, setGEnv, updGEnv,
12
13 LocalEnv(..),
14 readLEnv, setLEnv, updLEnv,
15
16 lookupTyCon,
17 lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
18
19 lookupInst, lookupFamInst
20 ) where
21
22 #include "HsVersions.h"
23
24 import HscTypes
25 import CoreSyn
26 import Class
27 import TyCon
28 import Type
29 import Var
30 import VarEnv
31 import Id
32 import Name
33 import NameEnv
34
35 import DsMonad
36 import PrelNames
37
38 import InstEnv
39 import FamInstEnv
40
41 import Panic
42 import Outputable
43 import FastString
44
45 -- ----------------------------------------------------------------------------
46 -- Vectorisation monad
47
48 data Builtins = Builtins {
49 parrayTyCon :: TyCon
50 , paClass :: Class
51 , closureTyCon :: TyCon
52 , mkClosureVar :: Var
53 , applyClosureVar :: Var
54 , mkClosurePVar :: Var
55 , applyClosurePVar :: Var
56 , lengthPAVar :: Var
57 , replicatePAVar :: Var
58 }
59
60 paDictTyCon :: Builtins -> TyCon
61 paDictTyCon = classTyCon . paClass
62
63 initBuiltins :: DsM Builtins
64 initBuiltins
65 = do
66 parrayTyCon <- dsLookupTyCon parrayTyConName
67 paClass <- dsLookupClass paClassName
68 closureTyCon <- dsLookupTyCon closureTyConName
69
70 mkClosureVar <- dsLookupGlobalId mkClosureName
71 applyClosureVar <- dsLookupGlobalId applyClosureName
72 mkClosurePVar <- dsLookupGlobalId mkClosurePName
73 applyClosurePVar <- dsLookupGlobalId applyClosurePName
74 lengthPAVar <- dsLookupGlobalId lengthPAName
75 replicatePAVar <- dsLookupGlobalId replicatePAName
76
77 return $ Builtins {
78 parrayTyCon = parrayTyCon
79 , paClass = paClass
80 , closureTyCon = closureTyCon
81 , mkClosureVar = mkClosureVar
82 , applyClosureVar = applyClosureVar
83 , mkClosurePVar = mkClosurePVar
84 , applyClosurePVar = applyClosurePVar
85 , lengthPAVar = lengthPAVar
86 , replicatePAVar = replicatePAVar
87 }
88
89 data GlobalEnv = GlobalEnv {
90 -- Mapping from global variables to their vectorised versions.
91 --
92 global_vars :: VarEnv CoreExpr
93
94 -- Exported variables which have a vectorised version
95 --
96 , global_exported_vars :: VarEnv (Var, Var)
97
98 -- Mapping from TyCons to their vectorised versions.
99 -- TyCons which do not have to be vectorised are mapped to
100 -- themselves.
101 --
102 , global_tycons :: NameEnv TyCon
103
104 -- Mapping from TyCons to their PA dictionaries
105 --
106 , global_tycon_pa :: NameEnv CoreExpr
107
108 -- External package inst-env & home-package inst-env for class
109 -- instances
110 --
111 , global_inst_env :: (InstEnv, InstEnv)
112
113 -- External package inst-env & home-package inst-env for family
114 -- instances
115 --
116 , global_fam_inst_env :: FamInstEnvs
117 }
118
119 data LocalEnv = LocalEnv {
120 -- Mapping from local variables to their vectorised and
121 -- lifted versions
122 --
123 local_vars :: VarEnv (CoreExpr, CoreExpr)
124
125 -- Mapping from tyvars to their PA dictionaries
126 , local_tyvar_pa :: VarEnv CoreExpr
127 }
128
129
130 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
131 initGlobalEnv info instEnvs famInstEnvs
132 = GlobalEnv {
133 global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
134 , global_exported_vars = emptyVarEnv
135 , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
136 , global_tycon_pa = emptyNameEnv
137 , global_inst_env = instEnvs
138 , global_fam_inst_env = famInstEnvs
139 }
140
141 emptyLocalEnv = LocalEnv {
142 local_vars = emptyVarEnv
143 , local_tyvar_pa = emptyVarEnv
144 }
145
146 -- FIXME
147 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
148 updVectInfo env tyenv info
149 = info {
150 vectInfoCCVar = global_exported_vars env
151 , vectInfoCCTyCon = tc_env
152 }
153 where
154 tc_env = mkNameEnv [(tc_name, (tc,tc'))
155 | tc <- typeEnvTyCons tyenv
156 , let tc_name = tyConName tc
157 , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
158
159 data VResult a = Yes GlobalEnv LocalEnv a | No
160
161 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
162
163 instance Monad VM where
164 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
165 VM p >>= f = VM $ \bi genv lenv -> do
166 r <- p bi genv lenv
167 case r of
168 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
169 No -> return No
170
171 noV :: VM a
172 noV = VM $ \_ _ _ -> return No
173
174 tryV :: VM a -> VM (Maybe a)
175 tryV (VM p) = VM $ \bi genv lenv ->
176 do
177 r <- p bi genv lenv
178 case r of
179 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
180 No -> return (Yes genv lenv Nothing)
181
182 maybeV :: VM (Maybe a) -> VM a
183 maybeV p = maybe noV return =<< p
184
185 orElseV :: VM a -> VM a -> VM a
186 orElseV p q = maybe q return =<< tryV p
187
188 localV :: VM a -> VM a
189 localV p = do
190 env <- readLEnv id
191 x <- p
192 setLEnv env
193 return x
194
195 closedV :: VM a -> VM a
196 closedV p = do
197 env <- readLEnv id
198 setLEnv emptyLocalEnv
199 x <- p
200 setLEnv env
201 return x
202
203 liftDs :: DsM a -> VM a
204 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
205
206 builtin :: (Builtins -> a) -> VM a
207 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
208
209 readGEnv :: (GlobalEnv -> a) -> VM a
210 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
211
212 setGEnv :: GlobalEnv -> VM ()
213 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
214
215 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
216 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
217
218 readLEnv :: (LocalEnv -> a) -> VM a
219 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
220
221 setLEnv :: LocalEnv -> VM ()
222 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
223
224 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
225 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
226
227 getInstEnv :: VM (InstEnv, InstEnv)
228 getInstEnv = readGEnv global_inst_env
229
230 getFamInstEnv :: VM FamInstEnvs
231 getFamInstEnv = readGEnv global_fam_inst_env
232
233 newLocalVar :: FastString -> Type -> VM Var
234 newLocalVar fs ty
235 = do
236 u <- liftDs newUnique
237 return $ mkSysLocal fs u ty
238
239 newTyVar :: FastString -> Kind -> VM Var
240 newTyVar fs k
241 = do
242 u <- liftDs newUnique
243 return $ mkTyVar (mkSysTvName u fs) k
244
245 lookupTyCon :: TyCon -> VM (Maybe TyCon)
246 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
247
248 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
249 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
250
251 extendTyVarPA :: Var -> CoreExpr -> VM ()
252 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
253
254 deleteTyVarPA :: Var -> VM ()
255 deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
256
257 -- Look up the dfun of a class instance.
258 --
259 -- The match must be unique - ie, match exactly one instance - but the
260 -- type arguments used for matching may be more specific than those of
261 -- the class instance declaration. The found class instances must not have
262 -- any type variables in the instance context that do not appear in the
263 -- instances head (i.e., no flexi vars); for details for what this means,
264 -- see the docs at InstEnv.lookupInstEnv.
265 --
266 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
267 lookupInst cls tys
268 = do { instEnv <- getInstEnv
269 ; case lookupInstEnv instEnv cls tys of
270 ([(inst, inst_tys)], _)
271 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
272 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
273 (ppr $ mkTyConApp (classTyCon cls) tys)
274 where
275 inst_tys' = [ty | Right ty <- inst_tys]
276 noFlexiVar = all isRight inst_tys
277 _other -> noV
278 }
279 where
280 isRight (Left _) = False
281 isRight (Right _) = True
282
283 -- Look up the representation tycon of a family instance.
284 --
285 -- The match must be unique - ie, match exactly one instance - but the
286 -- type arguments used for matching may be more specific than those of
287 -- the family instance declaration.
288 --
289 -- Return the instance tycon and its type instance. For example, if we have
290 --
291 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
292 --
293 -- then we have a coercion (ie, type instance of family instance coercion)
294 --
295 -- :Co:R42T Int :: T [Int] ~ :R42T Int
296 --
297 -- which implies that :R42T was declared as 'data instance T [a]'.
298 --
299 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
300 lookupFamInst tycon tys
301 = ASSERT( isOpenTyCon tycon )
302 do { instEnv <- getFamInstEnv
303 ; case lookupFamInstEnv instEnv tycon tys of
304 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
305 _other ->
306 pprPanic "VectMonad.lookupFamInst: not found: "
307 (ppr $ mkTyConApp tycon tys)
308 }
309
310 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
311 initV hsc_env guts info p
312 = do
313 eps <- hscEPS hsc_env
314 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
315 let instEnvs = (eps_inst_env eps, mg_inst_env guts)
316
317 Just r <- initDs hsc_env (mg_module guts)
318 (mg_rdr_env guts)
319 (mg_types guts)
320 (go instEnvs famInstEnvs)
321 return r
322 where
323
324 go instEnvs famInstEnvs =
325 do
326 builtins <- initBuiltins
327 r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
328 emptyLocalEnv
329 case r of
330 Yes genv _ x -> return $ Just (new_info genv, x)
331 No -> return Nothing
332
333 new_info genv = updVectInfo genv (mg_types guts) info
334