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 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
156 setFamInstEnv l_fam_inst genv
157 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
159 (g_fam_inst, _) = global_fam_inst_env genv
161 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
162 extendTyConsEnv ps genv
163 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
165 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
166 extendPAFunsEnv ps genv
167 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
169 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
171 = genv { global_pr_funs = mkNameEnv ps }
173 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
174 setBoxedTyConsEnv ps genv
175 = genv { global_boxed_tycons = mkNameEnv ps }
177 emptyLocalEnv = LocalEnv {
178 local_vars = emptyVarEnv
180 , local_tyvar_pa = emptyVarEnv
181 , local_bind_name = FSLIT("fn")
185 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
186 updVectInfo env tyenv info
188 vectInfoVar = global_exported_vars env
189 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
190 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
191 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
194 mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
195 | from <- from_tyenv tyenv
196 , let name = getName from
197 , Just to <- [lookupNameEnv (from_env env) name]]
199 data VResult a = Yes GlobalEnv LocalEnv a | No
201 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
203 instance Monad VM where
204 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
205 VM p >>= f = VM $ \bi genv lenv -> do
208 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
212 noV = VM $ \_ _ _ -> return No
214 traceNoV :: String -> SDoc -> VM a
215 traceNoV s d = pprTrace s d noV
217 tryV :: VM a -> VM (Maybe a)
218 tryV (VM p) = VM $ \bi genv lenv ->
222 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
223 No -> return (Yes genv lenv Nothing)
225 maybeV :: VM (Maybe a) -> VM a
226 maybeV p = maybe noV return =<< p
228 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
229 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
231 orElseV :: VM a -> VM a -> VM a
232 orElseV p q = maybe q return =<< tryV p
234 fixV :: (a -> VM a) -> VM a
235 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
237 unYes (Yes _ _ x) = x
239 localV :: VM a -> VM a
246 closedV :: VM a -> VM a
249 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
254 liftDs :: DsM a -> VM a
255 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
257 builtin :: (Builtins -> a) -> VM a
258 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
260 builtins :: (a -> Builtins -> b) -> VM (a -> b)
261 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
263 readGEnv :: (GlobalEnv -> a) -> VM a
264 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
266 setGEnv :: GlobalEnv -> VM ()
267 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
269 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
270 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
272 readLEnv :: (LocalEnv -> a) -> VM a
273 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
275 setLEnv :: LocalEnv -> VM ()
276 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
278 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
279 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
281 getInstEnv :: VM (InstEnv, InstEnv)
282 getInstEnv = readGEnv global_inst_env
284 getFamInstEnv :: VM FamInstEnvs
285 getFamInstEnv = readGEnv global_fam_inst_env
287 getBindName :: VM FastString
288 getBindName = readLEnv local_bind_name
290 inBind :: Id -> VM a -> VM a
292 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
295 cloneName :: (OccName -> OccName) -> Name -> VM Name
296 cloneName mk_occ name = liftM make (liftDs newUnique)
298 occ_name = mk_occ (nameOccName name)
300 make u | isExternalName name = mkExternalName u (nameModule name)
303 | otherwise = mkSystemName u occ_name
305 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
308 name <- cloneName mk_occ (getName id)
309 let id' | isExportedId id = Id.mkExportedLocalId name ty
310 | otherwise = Id.mkLocalId name ty
313 cloneVar :: Var -> VM Var
314 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
316 newExportedVar :: OccName -> Type -> VM Var
317 newExportedVar occ_name ty
319 mod <- liftDs getModuleDs
320 u <- liftDs newUnique
322 let name = mkExternalName u mod occ_name noSrcSpan
324 return $ Id.mkExportedLocalId name ty
326 newLocalVar :: FastString -> Type -> VM Var
329 u <- liftDs newUnique
330 return $ mkSysLocal fs u ty
332 newDummyVar :: Type -> VM Var
333 newDummyVar = newLocalVar FSLIT("ds")
335 newTyVar :: FastString -> Kind -> VM Var
338 u <- liftDs newUnique
339 return $ mkTyVar (mkSysTvName u fs) k
341 defGlobalVar :: Var -> Var -> VM ()
342 defGlobalVar v v' = updGEnv $ \env ->
343 env { global_vars = extendVarEnv (global_vars env) v v'
344 , global_exported_vars = upd (global_exported_vars env)
347 upd env | isExportedId v = extendVarEnv env v (v, v')
350 lookupVar :: Var -> VM (Scope Var (Var, Var))
353 r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
355 Just e -> return (Local e)
356 Nothing -> liftM Global
357 $ traceMaybeV "lookupVar" (ppr v)
358 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
360 lookupTyCon :: TyCon -> VM (Maybe TyCon)
362 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
364 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
366 defTyCon :: TyCon -> TyCon -> VM ()
367 defTyCon tc tc' = updGEnv $ \env ->
368 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
370 lookupDataCon :: DataCon -> VM (Maybe DataCon)
371 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
373 defDataCon :: DataCon -> DataCon -> VM ()
374 defDataCon dc dc' = updGEnv $ \env ->
375 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
377 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
378 lookupPrimPArray = liftDs . primPArray
380 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
381 lookupPrimMethod tycon = liftDs . primMethod tycon
383 lookupTyConPA :: TyCon -> VM (Maybe Var)
384 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
386 defTyConPA :: TyCon -> Var -> VM ()
387 defTyConPA tc pa = updGEnv $ \env ->
388 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
390 defTyConPAs :: [(TyCon, Var)] -> VM ()
391 defTyConPAs ps = updGEnv $ \env ->
392 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
393 [(tyConName tc, pa) | (tc, pa) <- ps] }
395 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
396 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
398 lookupTyConPR :: TyCon -> VM (Maybe Var)
399 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
401 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
402 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
405 defLocalTyVar :: TyVar -> VM ()
406 defLocalTyVar tv = updLEnv $ \env ->
407 env { local_tyvars = tv : local_tyvars env
408 , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
411 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
412 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
413 env { local_tyvars = tv : local_tyvars env
414 , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
417 localTyVars :: VM [TyVar]
418 localTyVars = readLEnv (reverse . local_tyvars)
420 -- Look up the dfun of a class instance.
422 -- The match must be unique - ie, match exactly one instance - but the
423 -- type arguments used for matching may be more specific than those of
424 -- the class instance declaration. The found class instances must not have
425 -- any type variables in the instance context that do not appear in the
426 -- instances head (i.e., no flexi vars); for details for what this means,
427 -- see the docs at InstEnv.lookupInstEnv.
430 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
432 = do { instEnv <- getInstEnv
433 ; case lookupInstEnv instEnv cls tys of
434 ([(inst, inst_tys)], _)
435 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
436 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
437 (ppr $ mkTyConApp (classTyCon cls) tys)
439 inst_tys' = [ty | Right ty <- inst_tys]
440 noFlexiVar = all isRight inst_tys
441 _other -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
444 isRight (Left _) = False
445 isRight (Right _) = True
448 -- Look up the representation tycon of a family instance.
450 -- The match must be unique - ie, match exactly one instance - but the
451 -- type arguments used for matching may be more specific than those of
452 -- the family instance declaration.
454 -- Return the instance tycon and its type instance. For example, if we have
456 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
458 -- then we have a coercion (ie, type instance of family instance coercion)
460 -- :Co:R42T Int :: T [Int] ~ :R42T Int
462 -- which implies that :R42T was declared as 'data instance T [a]'.
464 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
465 lookupFamInst tycon tys
466 = ASSERT( isOpenTyCon tycon )
467 do { instEnv <- getFamInstEnv
468 ; case lookupFamInstEnv instEnv tycon tys of
469 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
471 pprPanic "VectMonad.lookupFamInst: not found: "
472 (ppr $ mkTyConApp tycon tys)
475 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
476 initV hsc_env guts info p
478 Just r <- initDs hsc_env (mg_module guts)
487 builtins <- initBuiltins
488 let builtin_tycons = initBuiltinTyCons builtins
489 builtin_pas <- initBuiltinPAs builtins
490 builtin_prs <- initBuiltinPRs builtins
491 builtin_boxed <- initBuiltinBoxedTyCons builtins
493 eps <- ioToIOEnv $ hscEPS hsc_env
494 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
495 instEnvs = (eps_inst_env eps, mg_inst_env guts)
497 let genv = extendTyConsEnv builtin_tycons
498 . extendPAFunsEnv builtin_pas
499 . setPRFunsEnv builtin_prs
500 . setBoxedTyConsEnv builtin_boxed
501 $ initGlobalEnv info instEnvs famInstEnvs
503 r <- runVM p builtins genv emptyLocalEnv
505 Yes genv _ x -> return $ Just (new_info genv, x)
508 new_info genv = updVectInfo genv (mg_types guts) info