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
12 noV, tryV, maybeV, traceMaybeV, orElseV, fixV, localV, closedV, initV,
14 cloneName, cloneId, cloneVar,
15 newExportedVar, newLocalVar, newDummyVar, newTyVar,
17 Builtins(..), sumTyCon, prodTyCon, combinePAVar,
22 readGEnv, setGEnv, updGEnv,
25 readLEnv, setLEnv, updLEnv,
29 lookupVar, defGlobalVar,
30 lookupTyCon, defTyCon,
31 lookupDataCon, defDataCon,
32 lookupTyConPA, defTyConPA, defTyConPAs,
35 lookupPrimMethod, lookupPrimPArray,
36 lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
38 {-lookupInst,-} lookupFamInst
41 #include "HsVersions.h"
57 import TysPrim ( intPrimTy )
60 import IOEnv ( ioToIOEnv )
71 import SrcLoc ( noSrcSpan )
73 import Control.Monad ( liftM, zipWithM )
75 data Scope a b = Global a | Local b
77 -- ----------------------------------------------------------------------------
78 -- Vectorisation monad
80 data GlobalEnv = GlobalEnv {
81 -- Mapping from global variables to their vectorised versions.
83 global_vars :: VarEnv Var
85 -- Exported variables which have a vectorised version
87 , global_exported_vars :: VarEnv (Var, Var)
89 -- Mapping from TyCons to their vectorised versions.
90 -- TyCons which do not have to be vectorised are mapped to
93 , global_tycons :: NameEnv TyCon
95 -- Mapping from DataCons to their vectorised versions
97 , global_datacons :: NameEnv DataCon
99 -- Mapping from TyCons to their PA dfuns
101 , global_pa_funs :: NameEnv Var
103 -- Mapping from TyCons to their PR dfuns
104 , global_pr_funs :: NameEnv Var
106 -- Mapping from unboxed TyCons to their boxed versions
107 , global_boxed_tycons :: NameEnv TyCon
109 -- External package inst-env & home-package inst-env for class
112 , global_inst_env :: (InstEnv, InstEnv)
114 -- External package inst-env & home-package inst-env for family
117 , global_fam_inst_env :: FamInstEnvs
120 , global_bindings :: [(Var, CoreExpr)]
123 data LocalEnv = LocalEnv {
124 -- Mapping from local variables to their vectorised and
127 local_vars :: VarEnv (Var, Var)
129 -- In-scope type variables
131 , local_tyvars :: [TyVar]
133 -- Mapping from tyvars to their PA dictionaries
134 , local_tyvar_pa :: VarEnv CoreExpr
136 -- Local binding name
137 , local_bind_name :: FastString
140 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
141 initGlobalEnv info instEnvs famInstEnvs
143 global_vars = mapVarEnv snd $ vectInfoVar info
144 , global_exported_vars = emptyVarEnv
145 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
146 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
147 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
148 , global_pr_funs = emptyNameEnv
149 , global_boxed_tycons = emptyNameEnv
150 , global_inst_env = instEnvs
151 , global_fam_inst_env = famInstEnvs
152 , global_bindings = []
155 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
156 extendImportedVarsEnv ps genv
157 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
159 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
160 setFamInstEnv l_fam_inst genv
161 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
163 (g_fam_inst, _) = global_fam_inst_env genv
165 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
166 extendTyConsEnv ps genv
167 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
169 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
170 extendDataConsEnv ps genv
171 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
173 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
174 extendPAFunsEnv ps genv
175 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
177 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
179 = genv { global_pr_funs = mkNameEnv ps }
181 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
182 setBoxedTyConsEnv ps genv
183 = genv { global_boxed_tycons = mkNameEnv ps }
185 emptyLocalEnv = LocalEnv {
186 local_vars = emptyVarEnv
188 , local_tyvar_pa = emptyVarEnv
189 , local_bind_name = FSLIT("fn")
193 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
194 updVectInfo env tyenv info
196 vectInfoVar = global_exported_vars env
197 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
198 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
199 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
202 mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
203 | from <- from_tyenv tyenv
204 , let name = getName from
205 , Just to <- [lookupNameEnv (from_env env) name]]
207 data VResult a = Yes GlobalEnv LocalEnv a | No
209 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
211 instance Monad VM where
212 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
213 VM p >>= f = VM $ \bi genv lenv -> do
216 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
220 noV = VM $ \_ _ _ -> return No
222 traceNoV :: String -> SDoc -> VM a
223 traceNoV s d = pprTrace s d noV
225 tryV :: VM a -> VM (Maybe a)
226 tryV (VM p) = VM $ \bi genv lenv ->
230 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
231 No -> return (Yes genv lenv Nothing)
233 maybeV :: VM (Maybe a) -> VM a
234 maybeV p = maybe noV return =<< p
236 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
237 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
239 orElseV :: VM a -> VM a -> VM a
240 orElseV p q = maybe q return =<< tryV p
242 fixV :: (a -> VM a) -> VM a
243 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
245 unYes (Yes _ _ x) = x
247 localV :: VM a -> VM a
254 closedV :: VM a -> VM a
257 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
262 liftDs :: DsM a -> VM a
263 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
265 builtin :: (Builtins -> a) -> VM a
266 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
268 builtins :: (a -> Builtins -> b) -> VM (a -> b)
269 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
271 readGEnv :: (GlobalEnv -> a) -> VM a
272 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
274 setGEnv :: GlobalEnv -> VM ()
275 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
277 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
278 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
280 readLEnv :: (LocalEnv -> a) -> VM a
281 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
283 setLEnv :: LocalEnv -> VM ()
284 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
286 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
287 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
289 getInstEnv :: VM (InstEnv, InstEnv)
290 getInstEnv = readGEnv global_inst_env
292 getFamInstEnv :: VM FamInstEnvs
293 getFamInstEnv = readGEnv global_fam_inst_env
295 getBindName :: VM FastString
296 getBindName = readLEnv local_bind_name
298 inBind :: Id -> VM a -> VM a
300 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
303 cloneName :: (OccName -> OccName) -> Name -> VM Name
304 cloneName mk_occ name = liftM make (liftDs newUnique)
306 occ_name = mk_occ (nameOccName name)
308 make u | isExternalName name = mkExternalName u (nameModule name)
311 | otherwise = mkSystemName u occ_name
313 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
316 name <- cloneName mk_occ (getName id)
317 let id' | isExportedId id = Id.mkExportedLocalId name ty
318 | otherwise = Id.mkLocalId name ty
321 cloneVar :: Var -> VM Var
322 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
324 newExportedVar :: OccName -> Type -> VM Var
325 newExportedVar occ_name ty
327 mod <- liftDs getModuleDs
328 u <- liftDs newUnique
330 let name = mkExternalName u mod occ_name noSrcSpan
332 return $ Id.mkExportedLocalId name ty
334 newLocalVar :: FastString -> Type -> VM Var
337 u <- liftDs newUnique
338 return $ mkSysLocal fs u ty
340 newDummyVar :: Type -> VM Var
341 newDummyVar = newLocalVar FSLIT("ds")
343 newTyVar :: FastString -> Kind -> VM Var
346 u <- liftDs newUnique
347 return $ mkTyVar (mkSysTvName u fs) k
349 defGlobalVar :: Var -> Var -> VM ()
350 defGlobalVar v v' = updGEnv $ \env ->
351 env { global_vars = extendVarEnv (global_vars env) v v'
352 , global_exported_vars = upd (global_exported_vars env)
355 upd env | isExportedId v = extendVarEnv env v (v, v')
358 lookupVar :: Var -> VM (Scope Var (Var, Var))
361 r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
363 Just e -> return (Local e)
364 Nothing -> liftM Global
365 $ traceMaybeV "lookupVar" (ppr v)
366 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
368 lookupTyCon :: TyCon -> VM (Maybe TyCon)
370 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
372 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
374 defTyCon :: TyCon -> TyCon -> VM ()
375 defTyCon tc tc' = updGEnv $ \env ->
376 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
378 lookupDataCon :: DataCon -> VM (Maybe DataCon)
379 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
381 defDataCon :: DataCon -> DataCon -> VM ()
382 defDataCon dc dc' = updGEnv $ \env ->
383 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
385 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
386 lookupPrimPArray = liftDs . primPArray
388 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
389 lookupPrimMethod tycon = liftDs . primMethod tycon
391 lookupTyConPA :: TyCon -> VM (Maybe Var)
392 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
394 defTyConPA :: TyCon -> Var -> VM ()
395 defTyConPA tc pa = updGEnv $ \env ->
396 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
398 defTyConPAs :: [(TyCon, Var)] -> VM ()
399 defTyConPAs ps = updGEnv $ \env ->
400 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
401 [(tyConName tc, pa) | (tc, pa) <- ps] }
403 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
404 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
406 lookupTyConPR :: TyCon -> VM (Maybe Var)
407 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
409 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
410 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
413 defLocalTyVar :: TyVar -> VM ()
414 defLocalTyVar tv = updLEnv $ \env ->
415 env { local_tyvars = tv : local_tyvars env
416 , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
419 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
420 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
421 env { local_tyvars = tv : local_tyvars env
422 , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
425 localTyVars :: VM [TyVar]
426 localTyVars = readLEnv (reverse . local_tyvars)
428 -- Look up the dfun of a class instance.
430 -- The match must be unique - ie, match exactly one instance - but the
431 -- type arguments used for matching may be more specific than those of
432 -- the class instance declaration. The found class instances must not have
433 -- any type variables in the instance context that do not appear in the
434 -- instances head (i.e., no flexi vars); for details for what this means,
435 -- see the docs at InstEnv.lookupInstEnv.
438 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
440 = do { instEnv <- getInstEnv
441 ; case lookupInstEnv instEnv cls tys of
442 ([(inst, inst_tys)], _)
443 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
444 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
445 (ppr $ mkTyConApp (classTyCon cls) tys)
447 inst_tys' = [ty | Right ty <- inst_tys]
448 noFlexiVar = all isRight inst_tys
449 _other -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
452 isRight (Left _) = False
453 isRight (Right _) = True
456 -- Look up the representation tycon of a family instance.
458 -- The match must be unique - ie, match exactly one instance - but the
459 -- type arguments used for matching may be more specific than those of
460 -- the family instance declaration.
462 -- Return the instance tycon and its type instance. For example, if we have
464 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
466 -- then we have a coercion (ie, type instance of family instance coercion)
468 -- :Co:R42T Int :: T [Int] ~ :R42T Int
470 -- which implies that :R42T was declared as 'data instance T [a]'.
472 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
473 lookupFamInst tycon tys
474 = ASSERT( isOpenTyCon tycon )
475 do { instEnv <- getFamInstEnv
476 ; case lookupFamInstEnv instEnv tycon tys of
477 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
479 pprPanic "VectMonad.lookupFamInst: not found: "
480 (ppr $ mkTyConApp tycon tys)
483 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
484 initV hsc_env guts info p
486 Just r <- initDs hsc_env (mg_module guts)
495 builtins <- initBuiltins
496 builtin_vars <- initBuiltinVars builtins
497 builtin_tycons <- initBuiltinTyCons builtins
498 let builtin_datacons = initBuiltinDataCons builtins
499 builtin_pas <- initBuiltinPAs builtins
500 builtin_prs <- initBuiltinPRs builtins
501 builtin_boxed <- initBuiltinBoxedTyCons builtins
503 eps <- ioToIOEnv $ hscEPS hsc_env
504 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
505 instEnvs = (eps_inst_env eps, mg_inst_env guts)
507 let genv = extendImportedVarsEnv builtin_vars
508 . extendTyConsEnv builtin_tycons
509 . extendDataConsEnv builtin_datacons
510 . extendPAFunsEnv builtin_pas
511 . setPRFunsEnv builtin_prs
512 . setBoxedTyConsEnv builtin_boxed
513 $ initGlobalEnv info instEnvs famInstEnvs
515 r <- runVM p builtins genv emptyLocalEnv
517 Yes genv _ x -> return $ Just (new_info genv, x)
520 new_info genv = updVectInfo genv (mg_types guts) info