5 noV, tryV, maybeV, orElseV, localV, closedV, initV,
6 cloneName, newLocalVar, newTyVar,
8 Builtins(..), paDictTyCon,
12 readGEnv, setGEnv, updGEnv,
15 readLEnv, setLEnv, updLEnv,
17 defGlobalVar, lookupVar,
19 lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
21 lookupInst, lookupFamInst
24 #include "HsVersions.h"
47 import Control.Monad ( liftM )
49 data Scope a b = Global a | Local b
51 -- ----------------------------------------------------------------------------
52 -- Vectorisation monad
54 data Builtins = Builtins {
57 , closureTyCon :: TyCon
59 , applyClosureVar :: Var
60 , mkClosurePVar :: Var
61 , applyClosurePVar :: Var
63 , replicatePAVar :: Var
66 paDictTyCon :: Builtins -> TyCon
67 paDictTyCon = classTyCon . paClass
69 initBuiltins :: DsM Builtins
72 parrayTyCon <- dsLookupTyCon parrayTyConName
73 paClass <- dsLookupClass paClassName
74 closureTyCon <- dsLookupTyCon closureTyConName
76 mkClosureVar <- dsLookupGlobalId mkClosureName
77 applyClosureVar <- dsLookupGlobalId applyClosureName
78 mkClosurePVar <- dsLookupGlobalId mkClosurePName
79 applyClosurePVar <- dsLookupGlobalId applyClosurePName
80 lengthPAVar <- dsLookupGlobalId lengthPAName
81 replicatePAVar <- dsLookupGlobalId replicatePAName
84 parrayTyCon = parrayTyCon
86 , closureTyCon = closureTyCon
87 , mkClosureVar = mkClosureVar
88 , applyClosureVar = applyClosureVar
89 , mkClosurePVar = mkClosurePVar
90 , applyClosurePVar = applyClosurePVar
91 , lengthPAVar = lengthPAVar
92 , replicatePAVar = replicatePAVar
95 data GlobalEnv = GlobalEnv {
96 -- Mapping from global variables to their vectorised versions.
98 global_vars :: VarEnv CoreExpr
100 -- Exported variables which have a vectorised version
102 , global_exported_vars :: VarEnv (Var, Var)
104 -- Mapping from TyCons to their vectorised versions.
105 -- TyCons which do not have to be vectorised are mapped to
108 , global_tycons :: NameEnv TyCon
110 -- Mapping from TyCons to their PA dictionaries
112 , global_tycon_pa :: NameEnv CoreExpr
114 -- External package inst-env & home-package inst-env for class
117 , global_inst_env :: (InstEnv, InstEnv)
119 -- External package inst-env & home-package inst-env for family
122 , global_fam_inst_env :: FamInstEnvs
125 data LocalEnv = LocalEnv {
126 -- Mapping from local variables to their vectorised and
129 local_vars :: VarEnv (CoreExpr, CoreExpr)
131 -- Mapping from tyvars to their PA dictionaries
132 , local_tyvar_pa :: VarEnv CoreExpr
135 , local_bindings :: [(Var, CoreExpr)]
139 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
140 initGlobalEnv info instEnvs famInstEnvs
142 global_vars = mapVarEnv (Var . snd) $ vectInfoVar info
143 , global_exported_vars = emptyVarEnv
144 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
145 , global_tycon_pa = emptyNameEnv
146 , global_inst_env = instEnvs
147 , global_fam_inst_env = famInstEnvs
150 emptyLocalEnv = LocalEnv {
151 local_vars = emptyVarEnv
152 , local_tyvar_pa = emptyVarEnv
153 , local_bindings = []
157 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
158 updVectInfo env tyenv info
160 vectInfoVar = global_exported_vars env
161 , vectInfoTyCon = tc_env
164 tc_env = mkNameEnv [(tc_name, (tc,tc'))
165 | tc <- typeEnvTyCons tyenv
166 , let tc_name = tyConName tc
167 , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
169 data VResult a = Yes GlobalEnv LocalEnv a | No
171 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
173 instance Monad VM where
174 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
175 VM p >>= f = VM $ \bi genv lenv -> do
178 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
182 noV = VM $ \_ _ _ -> return No
184 tryV :: VM a -> VM (Maybe a)
185 tryV (VM p) = VM $ \bi genv lenv ->
189 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
190 No -> return (Yes genv lenv Nothing)
192 maybeV :: VM (Maybe a) -> VM a
193 maybeV p = maybe noV return =<< p
195 orElseV :: VM a -> VM a -> VM a
196 orElseV p q = maybe q return =<< tryV p
198 localV :: VM a -> VM a
205 closedV :: VM a -> VM a
208 setLEnv emptyLocalEnv
213 liftDs :: DsM a -> VM a
214 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
216 builtin :: (Builtins -> a) -> VM a
217 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
219 readGEnv :: (GlobalEnv -> a) -> VM a
220 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
222 setGEnv :: GlobalEnv -> VM ()
223 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
225 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
226 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
228 readLEnv :: (LocalEnv -> a) -> VM a
229 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
231 setLEnv :: LocalEnv -> VM ()
232 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
234 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
235 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
237 getInstEnv :: VM (InstEnv, InstEnv)
238 getInstEnv = readGEnv global_inst_env
240 getFamInstEnv :: VM FamInstEnvs
241 getFamInstEnv = readGEnv global_fam_inst_env
243 newLocalVar :: FastString -> Type -> VM Var
246 u <- liftDs newUnique
247 return $ mkSysLocal fs u ty
249 newTyVar :: FastString -> Kind -> VM Var
252 u <- liftDs newUnique
253 return $ mkTyVar (mkSysTvName u fs) k
255 defGlobalVar :: Var -> CoreExpr -> VM ()
256 defGlobalVar v e = updGEnv $ \env -> env { global_vars = extendVarEnv (global_vars env) v e }
258 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
261 r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
263 Just e -> return (Local e)
264 Nothing -> liftM Global
265 $ maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
267 lookupTyCon :: TyCon -> VM (Maybe TyCon)
268 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
270 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
271 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
273 extendTyVarPA :: Var -> CoreExpr -> VM ()
274 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
276 deleteTyVarPA :: Var -> VM ()
277 deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
279 -- Look up the dfun of a class instance.
281 -- The match must be unique - ie, match exactly one instance - but the
282 -- type arguments used for matching may be more specific than those of
283 -- the class instance declaration. The found class instances must not have
284 -- any type variables in the instance context that do not appear in the
285 -- instances head (i.e., no flexi vars); for details for what this means,
286 -- see the docs at InstEnv.lookupInstEnv.
288 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
290 = do { instEnv <- getInstEnv
291 ; case lookupInstEnv instEnv cls tys of
292 ([(inst, inst_tys)], _)
293 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
294 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
295 (ppr $ mkTyConApp (classTyCon cls) tys)
297 inst_tys' = [ty | Right ty <- inst_tys]
298 noFlexiVar = all isRight inst_tys
302 isRight (Left _) = False
303 isRight (Right _) = True
305 -- Look up the representation tycon of a family instance.
307 -- The match must be unique - ie, match exactly one instance - but the
308 -- type arguments used for matching may be more specific than those of
309 -- the family instance declaration.
311 -- Return the instance tycon and its type instance. For example, if we have
313 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
315 -- then we have a coercion (ie, type instance of family instance coercion)
317 -- :Co:R42T Int :: T [Int] ~ :R42T Int
319 -- which implies that :R42T was declared as 'data instance T [a]'.
321 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
322 lookupFamInst tycon tys
323 = ASSERT( isOpenTyCon tycon )
324 do { instEnv <- getFamInstEnv
325 ; case lookupFamInstEnv instEnv tycon tys of
326 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
328 pprPanic "VectMonad.lookupFamInst: not found: "
329 (ppr $ mkTyConApp tycon tys)
332 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
333 initV hsc_env guts info p
335 eps <- hscEPS hsc_env
336 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
337 let instEnvs = (eps_inst_env eps, mg_inst_env guts)
339 Just r <- initDs hsc_env (mg_module guts)
342 (go instEnvs famInstEnvs)
346 go instEnvs famInstEnvs =
348 builtins <- initBuiltins
349 r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
352 Yes genv _ x -> return $ Just (new_info genv, x)
355 new_info genv = updVectInfo genv (mg_types guts) info