module VectMonad (
+ Scope(..),
VM,
noV, tryV, maybeV, orElseV, localV, closedV, initV,
- newLocalVar, newTyVar,
+ cloneName, newLocalVar, newTyVar,
Builtins(..), paDictTyCon,
builtin,
LocalEnv(..),
readLEnv, setLEnv, updLEnv,
+ defGlobalVar, lookupVar,
lookupTyCon,
lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
import Var
import VarEnv
import Id
+import OccName
import Name
import NameEnv
import Outputable
import FastString
+import Control.Monad ( liftM )
+
+data Scope a b = Global a | Local b
+
-- ----------------------------------------------------------------------------
-- Vectorisation monad
-- instances
--
, global_fam_inst_env :: FamInstEnvs
+
+ -- Hoisted bindings
+ , global_bindings :: [(Var, CoreExpr)]
}
data LocalEnv = LocalEnv {
-- Mapping from tyvars to their PA dictionaries
, local_tyvar_pa :: VarEnv CoreExpr
-
- -- Hoisted bindings
- , local_bindings :: [(Var, CoreExpr)]
}
initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
initGlobalEnv info instEnvs famInstEnvs
= GlobalEnv {
- global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
+ global_vars = mapVarEnv (Var . snd) $ vectInfoVar info
, global_exported_vars = emptyVarEnv
- , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
+ , global_tycons = mapNameEnv snd $ vectInfoTyCon info
, global_tycon_pa = emptyNameEnv
, global_inst_env = instEnvs
, global_fam_inst_env = famInstEnvs
+ , global_bindings = []
}
emptyLocalEnv = LocalEnv {
local_vars = emptyVarEnv
, local_tyvar_pa = emptyVarEnv
- , local_bindings = []
}
-- FIXME
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
= info {
- vectInfoCCVar = global_exported_vars env
- , vectInfoCCTyCon = tc_env
+ vectInfoVar = global_exported_vars env
+ , vectInfoTyCon = tc_env
}
where
tc_env = mkNameEnv [(tc_name, (tc,tc'))
getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env
+cloneName :: (OccName -> OccName) -> Name -> VM Name
+cloneName mk_occ name = liftM make (liftDs newUnique)
+ where
+ occ_name = mk_occ (nameOccName name)
+
+ make u | isExternalName name = mkExternalName u (nameModule name)
+ occ_name
+ (nameSrcSpan name)
+ | otherwise = mkSystemName u occ_name
+
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
= do
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
+defGlobalVar :: Var -> Var -> VM ()
+defGlobalVar v v' = updGEnv $ \env ->
+ env { global_vars = extendVarEnv (global_vars env) v (Var v')
+ , global_exported_vars = upd (global_exported_vars env)
+ }
+ where
+ upd env | isExportedId v = extendVarEnv env v (v, v')
+ | otherwise = env
+
+lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
+lookupVar v
+ = do
+ r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+ case r of
+ Just e -> return (Local e)
+ Nothing -> liftM Global
+ $ maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
+
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)