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,
34 lookupPrimMethod, lookupPrimPArray,
35 lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
37 {-lookupInst,-} lookupFamInst
40 #include "HsVersions.h"
56 import TysPrim ( intPrimTy )
59 import IOEnv ( ioToIOEnv )
70 import SrcLoc ( noSrcSpan )
72 import Control.Monad ( liftM, zipWithM )
74 data Scope a b = Global a | Local b
76 -- ----------------------------------------------------------------------------
77 -- Vectorisation monad
79 data GlobalEnv = GlobalEnv {
80 -- Mapping from global variables to their vectorised versions.
82 global_vars :: VarEnv Var
84 -- Exported variables which have a vectorised version
86 , global_exported_vars :: VarEnv (Var, Var)
88 -- Mapping from TyCons to their vectorised versions.
89 -- TyCons which do not have to be vectorised are mapped to
92 , global_tycons :: NameEnv TyCon
94 -- Mapping from DataCons to their vectorised versions
96 , global_datacons :: NameEnv DataCon
98 -- Mapping from TyCons to their PA dfuns
100 , global_pa_funs :: NameEnv Var
102 -- Mapping from TyCons to their PR dfuns
103 , global_pr_funs :: NameEnv Var
105 -- External package inst-env & home-package inst-env for class
108 , global_inst_env :: (InstEnv, InstEnv)
110 -- External package inst-env & home-package inst-env for family
113 , global_fam_inst_env :: FamInstEnvs
116 , global_bindings :: [(Var, CoreExpr)]
119 data LocalEnv = LocalEnv {
120 -- Mapping from local variables to their vectorised and
123 local_vars :: VarEnv (Var, Var)
125 -- In-scope type variables
127 , local_tyvars :: [TyVar]
129 -- Mapping from tyvars to their PA dictionaries
130 , local_tyvar_pa :: VarEnv CoreExpr
132 -- Local binding name
133 , local_bind_name :: FastString
136 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
137 initGlobalEnv info instEnvs famInstEnvs
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 = []
150 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
151 setFamInstEnv l_fam_inst genv
152 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
154 (g_fam_inst, _) = global_fam_inst_env genv
156 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
157 extendTyConsEnv ps genv
158 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
160 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
161 extendPAFunsEnv ps genv
162 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
164 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
166 = genv { global_pr_funs = mkNameEnv ps }
168 emptyLocalEnv = LocalEnv {
169 local_vars = emptyVarEnv
171 , local_tyvar_pa = emptyVarEnv
172 , local_bind_name = FSLIT("fn")
176 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
177 updVectInfo env tyenv 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
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]]
190 data VResult a = Yes GlobalEnv LocalEnv a | No
192 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
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
199 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
203 noV = VM $ \_ _ _ -> return No
205 traceNoV :: String -> SDoc -> VM a
206 traceNoV s d = pprTrace s d noV
208 tryV :: VM a -> VM (Maybe a)
209 tryV (VM p) = VM $ \bi genv lenv ->
213 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
214 No -> return (Yes genv lenv Nothing)
216 maybeV :: VM (Maybe a) -> VM a
217 maybeV p = maybe noV return =<< p
219 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
220 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
222 orElseV :: VM a -> VM a -> VM a
223 orElseV p q = maybe q return =<< tryV p
225 fixV :: (a -> VM a) -> VM a
226 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
228 unYes (Yes _ _ x) = x
230 localV :: VM a -> VM a
237 closedV :: VM a -> VM a
240 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
245 liftDs :: DsM a -> VM a
246 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
248 builtin :: (Builtins -> a) -> VM a
249 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
251 builtins :: (a -> Builtins -> b) -> VM (a -> b)
252 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
254 readGEnv :: (GlobalEnv -> a) -> VM a
255 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
257 setGEnv :: GlobalEnv -> VM ()
258 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
260 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
261 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
263 readLEnv :: (LocalEnv -> a) -> VM a
264 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
266 setLEnv :: LocalEnv -> VM ()
267 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
269 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
270 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
272 getInstEnv :: VM (InstEnv, InstEnv)
273 getInstEnv = readGEnv global_inst_env
275 getFamInstEnv :: VM FamInstEnvs
276 getFamInstEnv = readGEnv global_fam_inst_env
278 getBindName :: VM FastString
279 getBindName = readLEnv local_bind_name
281 inBind :: Id -> VM a -> VM a
283 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
286 cloneName :: (OccName -> OccName) -> Name -> VM Name
287 cloneName mk_occ name = liftM make (liftDs newUnique)
289 occ_name = mk_occ (nameOccName name)
291 make u | isExternalName name = mkExternalName u (nameModule name)
294 | otherwise = mkSystemName u occ_name
296 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
299 name <- cloneName mk_occ (getName id)
300 let id' | isExportedId id = Id.mkExportedLocalId name ty
301 | otherwise = Id.mkLocalId name ty
304 cloneVar :: Var -> VM Var
305 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
307 newExportedVar :: OccName -> Type -> VM Var
308 newExportedVar occ_name ty
310 mod <- liftDs getModuleDs
311 u <- liftDs newUnique
313 let name = mkExternalName u mod occ_name noSrcSpan
315 return $ Id.mkExportedLocalId name ty
317 newLocalVar :: FastString -> Type -> VM Var
320 u <- liftDs newUnique
321 return $ mkSysLocal fs u ty
323 newDummyVar :: Type -> VM Var
324 newDummyVar = newLocalVar FSLIT("ds")
326 newTyVar :: FastString -> Kind -> VM Var
329 u <- liftDs newUnique
330 return $ mkTyVar (mkSysTvName u fs) k
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)
338 upd env | isExportedId v = extendVarEnv env v (v, v')
341 lookupVar :: Var -> VM (Scope Var (Var, Var))
344 r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
346 Just e -> return (Local e)
347 Nothing -> liftM Global
348 $ traceMaybeV "lookupVar" (ppr v)
349 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
351 lookupTyCon :: TyCon -> VM (Maybe TyCon)
353 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
355 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
357 defTyCon :: TyCon -> TyCon -> VM ()
358 defTyCon tc tc' = updGEnv $ \env ->
359 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
361 lookupDataCon :: DataCon -> VM (Maybe DataCon)
362 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
364 defDataCon :: DataCon -> DataCon -> VM ()
365 defDataCon dc dc' = updGEnv $ \env ->
366 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
368 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
369 lookupPrimPArray = liftDs . primPArray
371 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
372 lookupPrimMethod tycon = liftDs . primMethod tycon
374 lookupTyConPA :: TyCon -> VM (Maybe Var)
375 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
377 defTyConPA :: TyCon -> Var -> VM ()
378 defTyConPA tc pa = updGEnv $ \env ->
379 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
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] }
386 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
387 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
389 lookupTyConPR :: TyCon -> VM (Maybe Var)
390 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
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
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
404 localTyVars :: VM [TyVar]
405 localTyVars = readLEnv (reverse . local_tyvars)
407 -- Look up the dfun of a class instance.
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.
417 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
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)
426 inst_tys' = [ty | Right ty <- inst_tys]
427 noFlexiVar = all isRight inst_tys
428 _other -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
431 isRight (Left _) = False
432 isRight (Right _) = True
435 -- Look up the representation tycon of a family instance.
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.
441 -- Return the instance tycon and its type instance. For example, if we have
443 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
445 -- then we have a coercion (ie, type instance of family instance coercion)
447 -- :Co:R42T Int :: T [Int] ~ :R42T Int
449 -- which implies that :R42T was declared as 'data instance T [a]'.
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)
458 pprPanic "VectMonad.lookupFamInst: not found: "
459 (ppr $ mkTyConApp tycon tys)
462 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
463 initV hsc_env guts info p
465 Just r <- initDs hsc_env (mg_module guts)
474 builtins <- initBuiltins
475 let builtin_tycons = initBuiltinTyCons builtins
476 builtin_pas <- initBuiltinPAs builtins
477 builtin_prs <- initBuiltinPRs builtins
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)
483 let genv = extendTyConsEnv builtin_tycons
484 . extendPAFunsEnv builtin_pas
485 . setPRFunsEnv builtin_prs
486 $ initGlobalEnv info instEnvs famInstEnvs
488 r <- runVM p builtins genv emptyLocalEnv
490 Yes genv _ x -> return $ Just (new_info genv, x)
493 new_info genv = updVectInfo genv (mg_types guts) info