1bd450e237331ced4cade4dc7f226569ea933597
[ghc.git] / compiler / vectorise / VectMonad.hs
1 {-# OPTIONS -w #-}
2 -- The above warning supression flag is a temporary kludge.
3 -- While working on this module you are encouraged to remove it and fix
4 -- any warnings in the module. See
5 -- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
6 -- for details
7
8 module VectMonad (
9 Scope(..),
10 VM,
11
12 noV, tryV, maybeV, traceMaybeV, orElseV, fixV, localV, closedV, initV,
13 liftDs,
14 cloneName, cloneId, cloneVar,
15 newExportedVar, newLocalVar, newDummyVar, newTyVar,
16
17 Builtins(..), sumTyCon, prodTyCon,
18 builtin, builtins,
19
20 GlobalEnv(..),
21 setFamInstEnv,
22 readGEnv, setGEnv, updGEnv,
23
24 LocalEnv(..),
25 readLEnv, setLEnv, updLEnv,
26
27 getBindName, inBind,
28
29 lookupVar, defGlobalVar,
30 lookupTyCon, defTyCon,
31 lookupDataCon, defDataCon,
32 lookupTyConPA, defTyConPA, defTyConPAs,
33 lookupTyConPR,
34 lookupPrimMethod, lookupPrimPArray,
35 lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
36
37 {-lookupInst,-} lookupFamInst
38 ) where
39
40 #include "HsVersions.h"
41
42 import VectBuiltIn
43
44 import HscTypes
45 import CoreSyn
46 import TyCon
47 import DataCon
48 import Type
49 import Class
50 import Var
51 import VarEnv
52 import Id
53 import OccName
54 import Name
55 import NameEnv
56 import TysPrim ( intPrimTy )
57 import Module
58 import IfaceEnv
59 import IOEnv ( ioToIOEnv )
60
61 import DsMonad
62 import PrelNames
63
64 import InstEnv
65 import FamInstEnv
66
67 import Panic
68 import Outputable
69 import FastString
70 import SrcLoc ( noSrcSpan )
71
72 import Control.Monad ( liftM, zipWithM )
73
74 data Scope a b = Global a | Local b
75
76 -- ----------------------------------------------------------------------------
77 -- Vectorisation monad
78
79 data GlobalEnv = GlobalEnv {
80 -- Mapping from global variables to their vectorised versions.
81 --
82 global_vars :: VarEnv Var
83
84 -- Exported variables which have a vectorised version
85 --
86 , global_exported_vars :: VarEnv (Var, Var)
87
88 -- Mapping from TyCons to their vectorised versions.
89 -- TyCons which do not have to be vectorised are mapped to
90 -- themselves.
91 --
92 , global_tycons :: NameEnv TyCon
93
94 -- Mapping from DataCons to their vectorised versions
95 --
96 , global_datacons :: NameEnv DataCon
97
98 -- Mapping from TyCons to their PA dfuns
99 --
100 , global_pa_funs :: NameEnv Var
101
102 -- Mapping from TyCons to their PR dfuns
103 , global_pr_funs :: NameEnv Var
104
105 -- External package inst-env & home-package inst-env for class
106 -- instances
107 --
108 , global_inst_env :: (InstEnv, InstEnv)
109
110 -- External package inst-env & home-package inst-env for family
111 -- instances
112 --
113 , global_fam_inst_env :: FamInstEnvs
114
115 -- Hoisted bindings
116 , global_bindings :: [(Var, CoreExpr)]
117 }
118
119 data LocalEnv = LocalEnv {
120 -- Mapping from local variables to their vectorised and
121 -- lifted versions
122 --
123 local_vars :: VarEnv (Var, Var)
124
125 -- In-scope type variables
126 --
127 , local_tyvars :: [TyVar]
128
129 -- Mapping from tyvars to their PA dictionaries
130 , local_tyvar_pa :: VarEnv CoreExpr
131
132 -- Local binding name
133 , local_bind_name :: FastString
134 }
135
136 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
137 initGlobalEnv info instEnvs famInstEnvs
138 = GlobalEnv {
139 global_vars = mapVarEnv snd $ vectInfoVar info
140 , global_exported_vars = emptyVarEnv
141 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
142 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
143 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
144 , global_pr_funs = emptyNameEnv
145 , global_inst_env = instEnvs
146 , global_fam_inst_env = famInstEnvs
147 , global_bindings = []
148 }
149
150 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
151 setFamInstEnv l_fam_inst genv
152 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
153 where
154 (g_fam_inst, _) = global_fam_inst_env genv
155
156 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
157 extendTyConsEnv ps genv
158 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
159
160 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
161 extendPAFunsEnv ps genv
162 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
163
164 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
165 setPRFunsEnv ps genv
166 = genv { global_pr_funs = mkNameEnv ps }
167
168 emptyLocalEnv = LocalEnv {
169 local_vars = emptyVarEnv
170 , local_tyvars = []
171 , local_tyvar_pa = emptyVarEnv
172 , local_bind_name = FSLIT("fn")
173 }
174
175 -- FIXME
176 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
177 updVectInfo env tyenv info
178 = info {
179 vectInfoVar = global_exported_vars env
180 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
181 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
182 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
183 }
184 where
185 mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
186 | from <- from_tyenv tyenv
187 , let name = getName from
188 , Just to <- [lookupNameEnv (from_env env) name]]
189
190 data VResult a = Yes GlobalEnv LocalEnv a | No
191
192 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
193
194 instance Monad VM where
195 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
196 VM p >>= f = VM $ \bi genv lenv -> do
197 r <- p bi genv lenv
198 case r of
199 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
200 No -> return No
201
202 noV :: VM a
203 noV = VM $ \_ _ _ -> return No
204
205 traceNoV :: String -> SDoc -> VM a
206 traceNoV s d = pprTrace s d noV
207
208 tryV :: VM a -> VM (Maybe a)
209 tryV (VM p) = VM $ \bi genv lenv ->
210 do
211 r <- p bi genv lenv
212 case r of
213 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
214 No -> return (Yes genv lenv Nothing)
215
216 maybeV :: VM (Maybe a) -> VM a
217 maybeV p = maybe noV return =<< p
218
219 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
220 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
221
222 orElseV :: VM a -> VM a -> VM a
223 orElseV p q = maybe q return =<< tryV p
224
225 fixV :: (a -> VM a) -> VM a
226 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
227 where
228 unYes (Yes _ _ x) = x
229
230 localV :: VM a -> VM a
231 localV p = do
232 env <- readLEnv id
233 x <- p
234 setLEnv env
235 return x
236
237 closedV :: VM a -> VM a
238 closedV p = do
239 env <- readLEnv id
240 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
241 x <- p
242 setLEnv env
243 return x
244
245 liftDs :: DsM a -> VM a
246 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
247
248 builtin :: (Builtins -> a) -> VM a
249 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
250
251 builtins :: (a -> Builtins -> b) -> VM (a -> b)
252 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
253
254 readGEnv :: (GlobalEnv -> a) -> VM a
255 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
256
257 setGEnv :: GlobalEnv -> VM ()
258 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
259
260 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
261 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
262
263 readLEnv :: (LocalEnv -> a) -> VM a
264 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
265
266 setLEnv :: LocalEnv -> VM ()
267 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
268
269 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
270 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
271
272 getInstEnv :: VM (InstEnv, InstEnv)
273 getInstEnv = readGEnv global_inst_env
274
275 getFamInstEnv :: VM FamInstEnvs
276 getFamInstEnv = readGEnv global_fam_inst_env
277
278 getBindName :: VM FastString
279 getBindName = readLEnv local_bind_name
280
281 inBind :: Id -> VM a -> VM a
282 inBind id p
283 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
284 p
285
286 cloneName :: (OccName -> OccName) -> Name -> VM Name
287 cloneName mk_occ name = liftM make (liftDs newUnique)
288 where
289 occ_name = mk_occ (nameOccName name)
290
291 make u | isExternalName name = mkExternalName u (nameModule name)
292 occ_name
293 (nameSrcSpan name)
294 | otherwise = mkSystemName u occ_name
295
296 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
297 cloneId mk_occ id ty
298 = do
299 name <- cloneName mk_occ (getName id)
300 let id' | isExportedId id = Id.mkExportedLocalId name ty
301 | otherwise = Id.mkLocalId name ty
302 return id'
303
304 cloneVar :: Var -> VM Var
305 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
306
307 newExportedVar :: OccName -> Type -> VM Var
308 newExportedVar occ_name ty
309 = do
310 mod <- liftDs getModuleDs
311 u <- liftDs newUnique
312
313 let name = mkExternalName u mod occ_name noSrcSpan
314
315 return $ Id.mkExportedLocalId name ty
316
317 newLocalVar :: FastString -> Type -> VM Var
318 newLocalVar fs ty
319 = do
320 u <- liftDs newUnique
321 return $ mkSysLocal fs u ty
322
323 newDummyVar :: Type -> VM Var
324 newDummyVar = newLocalVar FSLIT("ds")
325
326 newTyVar :: FastString -> Kind -> VM Var
327 newTyVar fs k
328 = do
329 u <- liftDs newUnique
330 return $ mkTyVar (mkSysTvName u fs) k
331
332 defGlobalVar :: Var -> Var -> VM ()
333 defGlobalVar v v' = updGEnv $ \env ->
334 env { global_vars = extendVarEnv (global_vars env) v v'
335 , global_exported_vars = upd (global_exported_vars env)
336 }
337 where
338 upd env | isExportedId v = extendVarEnv env v (v, v')
339 | otherwise = env
340
341 lookupVar :: Var -> VM (Scope Var (Var, Var))
342 lookupVar v
343 = do
344 r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
345 case r of
346 Just e -> return (Local e)
347 Nothing -> liftM Global
348 $ traceMaybeV "lookupVar" (ppr v)
349 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
350
351 lookupTyCon :: TyCon -> VM (Maybe TyCon)
352 lookupTyCon tc
353 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
354
355 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
356
357 defTyCon :: TyCon -> TyCon -> VM ()
358 defTyCon tc tc' = updGEnv $ \env ->
359 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
360
361 lookupDataCon :: DataCon -> VM (Maybe DataCon)
362 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
363
364 defDataCon :: DataCon -> DataCon -> VM ()
365 defDataCon dc dc' = updGEnv $ \env ->
366 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
367
368 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
369 lookupPrimPArray = liftDs . primPArray
370
371 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
372 lookupPrimMethod tycon = liftDs . primMethod tycon
373
374 lookupTyConPA :: TyCon -> VM (Maybe Var)
375 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
376
377 defTyConPA :: TyCon -> Var -> VM ()
378 defTyConPA tc pa = updGEnv $ \env ->
379 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
380
381 defTyConPAs :: [(TyCon, Var)] -> VM ()
382 defTyConPAs ps = updGEnv $ \env ->
383 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
384 [(tyConName tc, pa) | (tc, pa) <- ps] }
385
386 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
387 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
388
389 lookupTyConPR :: TyCon -> VM (Maybe Var)
390 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
391
392 defLocalTyVar :: TyVar -> VM ()
393 defLocalTyVar tv = updLEnv $ \env ->
394 env { local_tyvars = tv : local_tyvars env
395 , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
396 }
397
398 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
399 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
400 env { local_tyvars = tv : local_tyvars env
401 , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
402 }
403
404 localTyVars :: VM [TyVar]
405 localTyVars = readLEnv (reverse . local_tyvars)
406
407 -- Look up the dfun of a class instance.
408 --
409 -- The match must be unique - ie, match exactly one instance - but the
410 -- type arguments used for matching may be more specific than those of
411 -- the class instance declaration. The found class instances must not have
412 -- any type variables in the instance context that do not appear in the
413 -- instances head (i.e., no flexi vars); for details for what this means,
414 -- see the docs at InstEnv.lookupInstEnv.
415 --
416 {-
417 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
418 lookupInst cls tys
419 = do { instEnv <- getInstEnv
420 ; case lookupInstEnv instEnv cls tys of
421 ([(inst, inst_tys)], _)
422 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
423 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
424 (ppr $ mkTyConApp (classTyCon cls) tys)
425 where
426 inst_tys' = [ty | Right ty <- inst_tys]
427 noFlexiVar = all isRight inst_tys
428 _other -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
429 }
430 where
431 isRight (Left _) = False
432 isRight (Right _) = True
433 -}
434
435 -- Look up the representation tycon of a family instance.
436 --
437 -- The match must be unique - ie, match exactly one instance - but the
438 -- type arguments used for matching may be more specific than those of
439 -- the family instance declaration.
440 --
441 -- Return the instance tycon and its type instance. For example, if we have
442 --
443 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
444 --
445 -- then we have a coercion (ie, type instance of family instance coercion)
446 --
447 -- :Co:R42T Int :: T [Int] ~ :R42T Int
448 --
449 -- which implies that :R42T was declared as 'data instance T [a]'.
450 --
451 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
452 lookupFamInst tycon tys
453 = ASSERT( isOpenTyCon tycon )
454 do { instEnv <- getFamInstEnv
455 ; case lookupFamInstEnv instEnv tycon tys of
456 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
457 _other ->
458 pprPanic "VectMonad.lookupFamInst: not found: "
459 (ppr $ mkTyConApp tycon tys)
460 }
461
462 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
463 initV hsc_env guts info p
464 = do
465 Just r <- initDs hsc_env (mg_module guts)
466 (mg_rdr_env guts)
467 (mg_types guts)
468 go
469 return r
470 where
471
472 go =
473 do
474 builtins <- initBuiltins
475 let builtin_tycons = initBuiltinTyCons builtins
476 builtin_pas <- initBuiltinPAs builtins
477 builtin_prs <- initBuiltinPRs builtins
478
479 eps <- ioToIOEnv $ hscEPS hsc_env
480 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
481 instEnvs = (eps_inst_env eps, mg_inst_env guts)
482
483 let genv = extendTyConsEnv builtin_tycons
484 . extendPAFunsEnv builtin_pas
485 . setPRFunsEnv builtin_prs
486 $ initGlobalEnv info instEnvs famInstEnvs
487
488 r <- runVM p builtins genv emptyLocalEnv
489 case r of
490 Yes genv _ x -> return $ Just (new_info genv, x)
491 No -> return Nothing
492
493 new_info genv = updVectInfo genv (mg_types guts) info
494